Function evaluate

train/probe_objbench_engine.py:93–144

(data_loader, model, device, args)

@torch.no_grad()
def evaluate(data_loader, model, device, args):
    if args.dataset in ['material', 'cloth']:
        criterion = torch.nn.CrossEntropyLoss()
    else:
        criterion = torch.nn.BCEWithLogitsLoss()
        sigmoid = torch.nn.Sigmoid()

    metric_logger = misc.MetricLogger(delimiter="  ")
    header = 'Test:'

    # switch to evaluation mode
    model.eval()

    for batch in metric_logger.log_every(data_loader, 40, header):
        images = batch[0]
        sensors = batch[1]
        target = batch[-1]
        images = images.to(device, non_blocking=True)
        sensors = sensors.to(device, non_blocking=True).int()
        target = target.to(device, non_blocking=True)

        if args.use_universal:
            sensors = torch.ones_like(sensors) * -1
            sensors = sensors.int()
        # compute output
        with torch.amp.autocast('cuda'):
            output = model(images, sensor_type = sensors)
            if args.dataset in ['rough', 'hard']:
                output = output.squeeze(1)
                target = target.float()
            loss = criterion(output, target)

        if args.dataset in ['material', 'cloth']:
            acc1, acc5 = accuracy(output, target, topk=(1,5))
        else:
            output = sigmoid(output)
            predictions = (output > 0.5).float()
            correct_predictions = (predictions == target).sum().item()
            acc1 = correct_predictions / target.size(0) * 100.0

        batch_size = images.shape[0]
        metric_logger.update(loss=loss.item())
        if args.dataset in ['material']:
            metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
        else:
            metric_logger.meters['acc1'].update(acc1, n=batch_size)

    metric_logger.synchronize_between_processes()
    print('* Acc@1 {top1.global_avg:.3f} loss {losses.global_avg:.3f}'
          .format(top1=metric_logger.acc1, losses=metric_logger.loss))

    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}