Source code for ezflow.functional.criterion.sequence

import torch
import torch.nn as nn

from ...config import configurable
from ..registry import FUNCTIONAL_REGISTRY


[docs]@FUNCTIONAL_REGISTRY.register() class SequenceLoss(nn.Module): """ Sequence loss for optical flow estimation. Used in **RAFT** (https://arxiv.org/abs/2003.12039) Parameters ---------- gamma : float Weight for the loss max_flow : float Maximum flow magnitude """ @configurable def __init__(self, gamma=0.8, max_flow=400, **kwargs): super(SequenceLoss, self).__init__() self.gamma = gamma self.max_flow = max_flow @classmethod def from_config(cls, cfg): return {"gamma": cfg.GAMMA, "max_flow": cfg.MAX_FLOW}
[docs] def forward(self, flow_preds, flow_gt, valid, **kwargs): # detect NaN nan_mask = (~torch.isnan(flow_gt)).float() flow_gt[torch.isnan(flow_gt)] = 0.0 n_preds = len(flow_preds) flow_loss = 0.0 valid = torch.squeeze(valid, dim=1) mag = torch.sqrt(torch.sum(flow_gt**2, dim=1)) valid = (valid >= 0.5) & (mag < self.max_flow) for i in range(n_preds): i_weight = self.gamma ** (n_preds - i - 1) i_loss = torch.abs(flow_preds[i] - flow_gt) flow_loss += i_weight * torch.mean((valid[:, None] * i_loss)) return flow_loss