Repository: Project-MONAI/MONAI
File: monai/losses/nacl_loss.py (Python)

Code before:
        self.cross_entropy = nn.CrossEntropyLoss()
        self.distance_type = distance_type
        self.alpha = alpha
        self.ks = kernel_size
        self.svls_layer: Any

        if kernel_ops == "mean":
            self.svls_layer = MeanFilter(spatial_dims=dim, size=kernel_size)
            self.svls_layer.filter = self.svls_layer.filter / (kernel_size**dim)
        if kernel_ops == "gaussian":
            self.svls_layer = GaussianFilter(spatial_dims=dim, sigma=sigma)

    def get_constr_target(self, mask: torch.Tensor) -> torch.Tensor:
        """
        Converts the mask to one hot represenation and applies the spatial filter.
        """
        rmask: torch.Tensor

        if self.dim == 2:
            oh_labels = F.one_hot(mask.to(torch.int64), num_classes=self.nc).contiguous().permute(0, 3, 1, 2).float()
            rmask = self.svls_layer(oh_labels)

        if self.dim == 3:
            oh_labels = F.one_hot(mask.to(torch.int64), num_classes=self.nc).contiguous().permute(0, 4, 1, 2, 3).float()
            rmask = self.svls_layer(oh_labels)

        return rmask

    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        loss_ce = self.cross_entropy(inputs, targets)

        utargets = self.get_constr_target(targets)

        if self.distance_type == "l1":
            loss_conf = utargets.sub(inputs).abs_().mean()
        elif self.distance_type == "l2":
            loss_conf = utargets.sub(inputs).pow_(2).abs_().mean()

        loss: torch.Tensor = loss_ce + self.alpha * loss_conf

        return loss


Review comment (performance):
Could you please add more description for the input shape for this function and also for the forward, such as the definition of dimension (we often assume all non-transform components assume input data shape [batch x channel x spatial_dims]) and also need to be one-hot or not. 

Here is an example in the `DiceCELoss`.
https://github.com/Project-MONAI/MONAI/blob/6c23fd06fc11667beedd0ba730d4104076a8db2d/monai/losses/dice.py#L776-L784

Code after:
        self.cross_entropy = nn.CrossEntropyLoss()
        self.distance_type = distance_type
        self.alpha = alpha
        self.ks = kernel_size
        self.svls_layer: Any

        if kernel_ops == "mean":
            self.svls_layer = MeanFilter(spatial_dims=dim, size=kernel_size)
            self.svls_layer.filter = self.svls_layer.filter / (kernel_size**dim)
        if kernel_ops == "gaussian":
            self.svls_layer = GaussianFilter(spatial_dims=dim, sigma=sigma)

    def get_constr_target(self, mask: torch.Tensor) -> torch.Tensor:
        """
        Converts the mask to one hot represenation and is smoothened with the selected spatial filter.

        Args:
            mask: the shape should be BH[WD].

        Returns:
            torch.Tensor: the shape would be BNH[WD], N being number of classes.
        """
        rmask: torch.Tensor

        if self.dim == 2:
            oh_labels = F.one_hot(mask.to(torch.int64), num_classes=self.nc).contiguous().permute(0, 3, 1, 2).float()
            rmask = self.svls_layer(oh_labels)

        if self.dim == 3:
            oh_labels = F.one_hot(mask.to(torch.int64), num_classes=self.nc).contiguous().permute(0, 4, 1, 2, 3).float()
            rmask = self.svls_layer(oh_labels)

        return rmask

    def forward(self, inputs: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        """
        Computes standard cross-entropy loss and constraints it neighbor aware logit penalty.

        Args:
            inputs: the shape should be BNH[WD], where N is the number of classes.
            targets: the shape should be BH[WD].

        Returns:
            torch.Tensor: value of the loss.

        Example:
            >>> import torch
            >>> from monai.losses import NACLLoss
            >>> B, N, H, W = 8, 3, 64, 64
            >>> input = torch.rand(B, N, H, W)
            >>> target = torch.randint(0, N, (B, H, W))