Function evaluate

Train/MatchingNet/evaluate.py:27–63

(model: nn.DataParallel[FlowFormerCov], loader: DataLoader, length: int, usewandb)

def evaluate(model: nn.DataParallel[FlowFormerCov], loader: DataLoader, length: int, usewandb) -> List[Dict]:
    model.eval()
    with torch.no_grad():
        step = 0
        metric_list = []
        for frameData in ColoredTqdm(loader, desc="Evaluation", total=length):

            img1, img2 = frameData.cur.imageL.cuda(), frameData.nxt.imageL.cuda()
            gt_flow = frameData.cur.gtFlow.cuda()
            # flow, cov = model.module.inference(img1, img2)
            flow_pre, cov_pre = model.forward(img1, img2)
            flow, cov = flow_pre[0], torch.exp(2 * cov_pre[0])
            
            error_mask = (gt_flow.norm(dim=1) < 240)
            flow_mask = error_mask.unsqueeze(1).expand_as(gt_flow) # for masking the flow
            
            MSE = (flow - gt_flow)**2
            EPE = (flow - gt_flow).norm(dim=1)
            masked_EPE = EPE[error_mask]
            
            cov_dist = cov.sqrt().norm(dim=1)
            cov_ratio = (cov_dist / EPE)
            eval_loss = MSE / (2 * cov) + 0.5 * torch.log(cov)
            step += 1
            if step > length:
                break
            metric = {
                'mse': MSE.mean().item(),
                'epe': EPE.mean().item(),
                'masked_epe': masked_EPE.mean().item(), # 'mask for epe < 200
                'cov_dist': cov_dist.mean().item(),
                'cov_ratio': cov_ratio.mean().item(),
                'eval_loss': eval_loss.mean().item()
            }

            metric_list.append(metric)
        return metric_list


if __name__ == "__main__":