INNER CODE UNIT · Python

forward

nianticlabs/simplerecon · losses.py:58

    def forward(self, normals_gt_b3hw: Tensor, normals_pred_b3hw: Tensor) -> Tensor:

        normals_mask_b1hw = torch.logical_and(
            normals_gt_b3hw.isfinite().all(dim=1, keepdim=True),
            normals_pred_b3hw.isfinite().all(dim=1, keepdim=True))

        normals_pred_b3hw = normals_pred_b3hw.masked_fill(~normals_mask_b1hw, 1.0)
        normals_gt_b3hw = normals_gt_b3hw.masked_fill(~normals_mask_b1hw, 1.0)

        with torch.cuda.amp.autocast(enabled=False):
            normals_dot_b1hw = 0.5 * (
                                        1.0 - torch.einsum(
                                                "bchw, bchw -> bhw", 
                                                normals_pred_b3hw, 
                                                normals_gt_b3hw,
                                            )
                                        ).unsqueeze(1)
        normals_loss = normals_dot_b1hw.masked_select(normals_mask_b1hw).mean()

View source record →