Class NetAdapter

monai/networks/nets/netadapter.py:26–123

Wrapper to replace the last layer of model by convolutional layer or FC layer. See also: :py:class:`monai.networks.nets.TorchVisionFCModel` Args: model: a PyTorch model, which can be both 2D and 3D models. typically, it can be a pretrained model in Torchvision, lik

class NetAdapter(torch.nn.Module):
    """
    Wrapper to replace the last layer of model by convolutional layer or FC layer.

    See also: :py:class:`monai.networks.nets.TorchVisionFCModel`

    Args:
        model: a PyTorch model, which can be both 2D and 3D models. typically, it can be a pretrained model
            in Torchvision, like: ``resnet18``, ``resnet34``, ``resnet50``, ``resnet101``, ``resnet152``, etc.
            more details: https://pytorch.org/vision/stable/models.html.
        num_classes: number of classes for the last classification layer. Default to 1.
        dim: number of supported spatial dimensions in the specified model, depends on the model implementation.
            default to 2 as most Torchvision models are for 2D image processing.
        in_channels: number of the input channels of last layer. if None, get it from `in_features` of last layer.
        use_conv: whether to use convolutional layer to replace the last layer, default to False.
        pool: parameters for the pooling layer, it should be a tuple, the first item is name of the pooling layer,
            the second item is dictionary of the initialization args. if None, will not replace the `layers[-2]`.
            default to `("avg", {"kernel_size": 7, "stride": 1})`.
        bias: the bias value when replacing the last layer. if False, the layer will not learn an additive bias,
            default to True.
        fc_name: the corresponding layer attribute of the last fully connected layer. Defaults to ``"fc"``.
        node_name: the corresponding feature extractor node name of `model`.
            Defaults to "", the extractor is not in use.

    """

    def __init__(
        self,
        model: torch.nn.Module,
        num_classes: int = 1,
        dim: int = 2,
        in_channels: int | None = None,
        use_conv: bool = False,
        pool: tuple[str, dict[str, Any]] | None = ("avg", {"kernel_size": 7, "stride": 1}),
        bias: bool = True,
        fc_name: str = "fc",
        node_name: str = "",
    ):
        super().__init__()
        layers = list(model.children())
        orig_fc = look_up_named_module(fc_name, model)
        if orig_fc is None:
            orig_fc = layers[-1]
        # guess the number of input channels of the last fully connected layer
        in_channels_: int
        if in_channels is None:
            if not hasattr(orig_fc, "in_features"):
                raise ValueError("please specify input channels of the last fully connected layer with `in_channels`.")
            in_channels_ = orig_fc.in_features

        else:
            in_channels_ = in_channels

        # modify the input model, depending on whether to replace the last pooling layer ``pool``
        if pool is None:  # no modification of pooling
            if node_name != "":
                raise ValueError("`node_name` is not compatible with `pool=None`, please set `pool=''`.")
            # we just drop the model's fully connected layer or set it to identity