INNER CODE UNIT · Python

si_loss

nianticlabs/simplerecon · losses.py:50

        si_loss = torch.sqrt(
            (log_diff ** 2).mean() - self.si_lambda * (log_diff.mean() ** 2)
        )

        return si_loss


class NormalsLoss(nn.Module):
    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):

View source record →