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

Code before:

    def __init__(
        self,
        spatial_dims: int,
        network_type: str = PercetualNetworkType.alex,
        is_fake_3d: bool = True,
        fake_3d_ratio: float = 0.5,
        cache_dir: str | None = None,
        pretrained: bool = True,
        pretrained_path: str | None = None,
        pretrained_state_dict_key: str | None = None,
        channelwise: bool = False,
    ):
        super().__init__()

        if spatial_dims not in [2, 3]:
            raise NotImplementedError("Perceptual loss is implemented only in 2D and 3D.")

        if (spatial_dims == 2 or is_fake_3d) and "medicalnet_" in network_type:
            raise ValueError(
                "MedicalNet networks are only compatible with ``spatial_dims=3``."
                "Argument is_fake_3d must be set to False."
            )

        if channelwise and "medicalnet_" not in network_type:
            raise ValueError("Channelwise loss is only compatible with MedicalNet networks.")

        if network_type.lower() not in list(PercetualNetworkType):
            raise ValueError(
                "Unrecognised criterion entered for Adversarial Loss. Must be one in: %s"
                % ", ".join(PercetualNetworkType)
            )

        if cache_dir:
            torch.hub.set_dir(cache_dir)
            # raise a warning that this may change the default cache dir for all torch.hub calls
            warnings.warn(
                f"Setting cache_dir to {cache_dir}, this may change the default cache dir for all torch.hub calls."
            )

        self.spatial_dims = spatial_dims
        self.perceptual_function: nn.Module
        if spatial_dims == 3 and is_fake_3d is False:
            self.perceptual_function = MedicalNetPerceptualSimilarity(
                net=network_type, verbose=False, channelwise=channelwise
            )
        elif "radimagenet_" in network_type:
            self.perceptual_function = RadImageNetPerceptualSimilarity(net=network_type, verbose=False)
        elif network_type == "resnet50":
            self.perceptual_function = TorchvisionModelPerceptualSimilarity(
                net=network_type,

Review comment (suggestion):
```suggestion
        if channel_wise and "medicalnet_" not in network_type:
            raise ValueError("Channel wise loss is only compatible with MedicalNet networks.")
```

Code after:
        channel_wise: if True, the loss is returned per channel. Otherwise the loss is averaged over the channels.
                Defaults to ``False``.
    """

    def __init__(
        self,
        spatial_dims: int,
        network_type: str = PercetualNetworkType.alex,
        is_fake_3d: bool = True,
        fake_3d_ratio: float = 0.5,
        cache_dir: str | None = None,
        pretrained: bool = True,
        pretrained_path: str | None = None,
        pretrained_state_dict_key: str | None = None,
        channel_wise: bool = False,
    ):
        super().__init__()

        if spatial_dims not in [2, 3]:
            raise NotImplementedError("Perceptual loss is implemented only in 2D and 3D.")

        if (spatial_dims == 2 or is_fake_3d) and "medicalnet_" in network_type:
            raise ValueError(
                "MedicalNet networks are only compatible with ``spatial_dims=3``."
                "Argument is_fake_3d must be set to False."
            )

        if channel_wise and "medicalnet_" not in network_type:
            raise ValueError("Channel-wise loss is only compatible with MedicalNet networks.")

        if network_type.lower() not in list(PercetualNetworkType):
            raise ValueError(
                "Unrecognised criterion entered for Adversarial Loss. Must be one in: %s"
                % ", ".join(PercetualNetworkType)
            )

        if cache_dir:
            torch.hub.set_dir(cache_dir)
            # raise a warning that this may change the default cache dir for all torch.hub calls
            warnings.warn(
                f"Setting cache_dir to {cache_dir}, this may change the default cache dir for all torch.hub calls."
            )

        self.spatial_dims = spatial_dims
        self.perceptual_function: nn.Module
        if spatial_dims == 3 and is_fake_3d is False:
            self.perceptual_function = MedicalNetPerceptualSimilarity(
                net=network_type, verbose=False, channel_wise=channel_wise
            )
        elif "radimagenet_" in network_type:
            self.perceptual_function = RadImageNetPerceptualSimilarity(net=network_type, verbose=False)