Repository: Project-MONAI/MONAI
File: monai/networks/layers/factories.py (Python)

Code before:
    `InstanceNorm3dNVFuser` is a faster version of InstanceNorm layer and implemented in `apex`.
    It only supports 3d tensors as the input. It also requires to use with CUDA and non-Windows OS.
    In this function, if the required library `apex.normalization.InstanceNorm3dNVFuser` does not exist,
    `nn.InstanceNorm3d` will be returned instead.
    This layer is based on a customized autograd function, which is not supported in TorchScript currently.
    Please switch to use `nn.InstanceNorm3d` if TorchScript is necessary.

    Please check the following link for more details about how to install `apex`:
    https://github.com/NVIDIA/apex#installation

    """

    if dim != 3:
        types = (nn.InstanceNorm1d, nn.InstanceNorm2d)
        warnings.warn(f"`InstanceNorm3dNVFuser` only supports 3d cases, use {types[dim - 1]} instead.")
        return types[dim - 1]

    if not has_nvfuser_instance_norm():
        warnings.warn(
            "`apex.normalization.InstanceNorm3dNVFuser` is not installed properly, use nn.InstanceNorm3d instead."
        )
        return nn.InstanceNorm3d
    return optional_import("apex.normalization", name="InstanceNorm3dNVFuser")[0]


Act.add("elu", lambda: nn.modules.ELU)
Act.add("relu", lambda: nn.modules.ReLU)
Act.add("leakyrelu", lambda: nn.modules.LeakyReLU)
Act.add("prelu", lambda: nn.modules.PReLU)
Act.add("relu6", lambda: nn.modules.ReLU6)
Act.add("selu", lambda: nn.modules.SELU)
Act.add("celu", lambda: nn.modules.CELU)
Act.add("gelu", lambda: nn.modules.GELU)
Act.add("sigmoid", lambda: nn.modules.Sigmoid)
Act.add("tanh", lambda: nn.modules.Tanh)
Act.add("softmax", lambda: nn.modules.Softmax)
Act.add("logsoftmax", lambda: nn.modules.LogSoftmax)


@Act.factory_item("swish")
def swish_factory():
    from monai.networks.blocks.activation import Swish

    return Swish


@Act.factory_item("memswish")
def memswish_factory():
    from monai.networks.blocks.activation import MemoryEfficientSwish

    return MemoryEfficientSwish

Review comment (refactor):
thanks for the refactoring! with the new API, I suspect the `lambda: nn.modules.ELU` could be further simplified as `nn.modules.ELU`...

Code after:


@Norm.factory_function("batch")
def batch_factory(dim: int) -> type[nn.BatchNorm1d | nn.BatchNorm2d | nn.BatchNorm3d]:
    """
    Batch normalization layers in 1,2,3 dimensions.

    Args:
        dim: desired dimension of the batch normalization layer

    Returns:
        BatchNorm[dim]d
    """
    types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)
    return types[dim - 1]


@Norm.factory_function("instance_nvfuser")
def instance_nvfuser_factory(dim):
    """
    `InstanceNorm3dNVFuser` is a faster version of InstanceNorm layer and implemented in `apex`.
    It only supports 3d tensors as the input. It also requires to use with CUDA and non-Windows OS.
    In this function, if the required library `apex.normalization.InstanceNorm3dNVFuser` does not exist,
    `nn.InstanceNorm3d` will be returned instead.
    This layer is based on a customized autograd function, which is not supported in TorchScript currently.
    Please switch to use `nn.InstanceNorm3d` if TorchScript is necessary.

    Please check the following link for more details about how to install `apex`:
    https://github.com/NVIDIA/apex#installation

    """

    if dim != 3:
        types = (nn.InstanceNorm1d, nn.InstanceNorm2d)
        warnings.warn(f"`InstanceNorm3dNVFuser` only supports 3d cases, use {types[dim - 1]} instead.")
        return types[dim - 1]

    if not has_nvfuser_instance_norm():
        warnings.warn(
            "`apex.normalization.InstanceNorm3dNVFuser` is not installed properly, use nn.InstanceNorm3d instead."
        )
        return nn.InstanceNorm3d
    return optional_import("apex.normalization", name="InstanceNorm3dNVFuser")[0]


Norm.add_factory_class("group", nn.GroupNorm)
Norm.add_factory_class("layer", nn.LayerNorm)
Norm.add_factory_class("localresponse", nn.LocalResponseNorm)
Norm.add_factory_class("syncbatch", nn.SyncBatchNorm)