Loss function defined over sequence of flow predictions
(flow_preds, flow_gt, valid, cfg)
| 3 | MAX_FLOW = 400 |
| 4 | |
| 5 | def sequence_loss(flow_preds, flow_gt, valid, cfg): |
| 6 | """ Loss function defined over sequence of flow predictions """ |
| 7 | |
| 8 | gamma = cfg.gamma |
| 9 | max_flow = cfg.max_flow |
| 10 | n_predictions = len(flow_preds) |
| 11 | flow_loss = 0.0 |
| 12 | flow_gt_thresholds = [5, 10, 20] |
| 13 | |
| 14 | # exlude invalid pixels and extremely large diplacements |
| 15 | mag = torch.sum(flow_gt**2, dim=1).sqrt() |
| 16 | valid = (valid >= 0.5) & (mag < max_flow) |
| 17 | |
| 18 | for i in range(n_predictions): |
| 19 | i_weight = gamma**(n_predictions - i - 1) |
| 20 | i_loss = (flow_preds[i] - flow_gt).abs() |
| 21 | flow_loss += i_weight * (valid[:, None] * i_loss).mean() |
| 22 | |
| 23 | epe = torch.sum((flow_preds[-1] - flow_gt)**2, dim=1).sqrt() |
| 24 | epe = epe.view(-1)[valid.view(-1)] |
| 25 | |
| 26 | metrics = { |
| 27 | 'epe': epe.mean().item(), |
| 28 | '1px': (epe < 1).float().mean().item(), |
| 29 | '3px': (epe < 3).float().mean().item(), |
| 30 | '5px': (epe < 5).float().mean().item(), |
| 31 | } |
| 32 | |
| 33 | flow_gt_length = torch.sum(flow_gt**2, dim=1).sqrt() |
| 34 | flow_gt_length = flow_gt_length.view(-1)[valid.view(-1)] |
| 35 | for t in flow_gt_thresholds: |
| 36 | e = epe[flow_gt_length < t] |
| 37 | metrics.update({ |
| 38 | f"{t}-th-5px": (e < 5).float().mean().item() |
| 39 | }) |
| 40 | |
| 41 | |
| 42 | return flow_loss, metrics |
| 43 |