Function evaluate

train/probe_touchd_engine.py:116–166

(data_loader, model, device, args, epoch)

@torch.no_grad()
def evaluate(data_loader, model, device, args, epoch):

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

    # switch to evaluation mode
    model.eval()
    all_predictions = []
    all_targets = []
    all_force_scales = []

    for batch in metric_logger.log_every(data_loader, 40, header):

        images = batch[0].to(device, non_blocking=True)
        sensors = batch[1].to(device, non_blocking=True).int()
        target = batch[2].to(device, non_blocking=True)
        force_scale = batch[3].to(device, non_blocking=True)

        # compute output
        with torch.amp.autocast('cuda'):
            output = model(images, sensor_type = sensors)

            all_predictions.append(output.detach())
            all_targets.append(target.detach())
            all_force_scales.append(force_scale.detach())


    all_predictions = torch.cat(all_predictions, dim=0)
    all_targets = torch.cat(all_targets, dim=0)
    all_force_scales = torch.cat(all_force_scales, dim=0)

    all_predictions = all_predictions * all_force_scales
    all_targets = all_targets * all_force_scales

    forces_rmse_xyz = torch.sqrt(((all_predictions - all_targets) ** 2).mean(dim=0)) * 1000  # in mN
    total_rmse = forces_rmse_xyz.sum()

    # print(total_rmse.shape, forces_rmse_xyz.shape)
    metric_logger.update(rmse_x=forces_rmse_xyz[0].item())
    metric_logger.update(rmse_y=forces_rmse_xyz[1].item())
    metric_logger.update(rmse_z=forces_rmse_xyz[2].item())
    metric_logger.update(rmse=total_rmse.item())

    # gather the stats from all processes
    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))
    print("Averaged stats:", metric_logger)
    plot_correlation(all_targets.cpu().numpy(), all_predictions.cpu().numpy(), args.log_dir, epoch)

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