INNER CODE UNIT · Python

MSGradientLoss

nianticlabs/simplerecon · losses.py:11

class MSGradientLoss(nn.Module):
    def __init__(self, num_scales: int = 4):
        super().__init__()

        self.num_scales = num_scales

    def forward(self, depth_gt: Tensor, depth_pred: Tensor) -> Tensor:

        # Create the gradient pyramids
        depth_pred_pyr = pyrdown(depth_pred, self.num_scales)
        depth_gtn_pyr = pyrdown(depth_gt, self.num_scales)

        grad_loss = torch.tensor(0, dtype=depth_gt.dtype, device=depth_gt.device)
        for depth_pred_down, depth_gtn_down in zip(depth_pred_pyr, depth_gtn_pyr):

            depth_gtn_grad = kornia.filters.spatial_gradient(depth_gtn_down)

            mask_down_b = depth_gtn_grad.isfinite().all(dim=1, keepdim=True)

View source record →