Make the flexible unet more extensible to different backbones

monai/networks/blocks/__init__.py

@@ -19,6 +19,7 @@
 from .dints_block import ActiConvNormBlock, FactorizedIncreaseBlock, FactorizedReduceBlock, P3DActiConvNormBlock
 from .downsample import MaxAvgPool
 from .dynunet_block import UnetBasicBlock, UnetOutBlock, UnetResBlock, UnetUpBlock, get_output_padding, get_padding
+from .encoder import BaseEncoder
 from .fcn import FCN, GCN, MCFCN, Refine
 from .feature_pyramid_network import ExtraFPNBlock, FeaturePyramidNetwork, LastLevelMaxPool, LastLevelP6P7
 from .localnet_block import LocalNetDownSampleBlock, LocalNetFeatureExtractorBlock, LocalNetUpSampleBlock

monai/networks/blocks/encoder.py

@@ -0,0 +1,84 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABCMeta, abstractmethod
+from typing import Dict, List, Tuple
+
+__all__ = ["BaseEncoder"]
+
+
+class BaseEncoder(metaclass=ABCMeta):
+    """
+    Abstract class defines interface of encoders in flexible unet.
+    Encoders in flexible unet must derive from this class. Each interface method
+    should return a list containing relative information about a series of newtworks
+    defined by encoder. For example, the efficient-net encoder implement 10 basic
+    network structures in one encoder. When calling `get_encoder_name_string_list`
+    function, a string list like ["efficientnet-b0", "efficientnet-b1" ... "efficientnet-l2"]
+    should be returned.
+    """
+
+    @classmethod
+    @abstractmethod
+    def get_encoder_parameters(cls) -> List[Dict]:
+        """
+        Get parameter list to initialize encoder networks.
+        Each parameter dict must have `spatial_dims`, `in_channels`
+        and `pretrained` parameters.
+        The reason that this function should return a list is that a
+        series of encoders can be implemented by one encoder class
+        given different initialization parameters. Each parameter dict
+        in return list should be able to initialize a unique encoder.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def num_channels_per_output(cls) -> List[Tuple[int, ...]]:
+        """
+        Get number of output features' channels.
+        The reason that this function should return a list is that a
+        series of encoders can be implemented by one encoder class
+        given different initialization parameters. And it is possible
+        that different encoders have different output feature map
+        channels. Therefore a list of output feature map channel tuples
+        corresponding to each encoder should be returned by this method.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def num_outputs(cls) -> List[int]:
+        """
+        Get number of outputs of encoder.
+        The reason that this function should return a list is that a
+        series of encoders can be implemented by one encoder class
+        given different initialization parameters. And it is possible
+        that different encoders have different output feature numbers.
+        Therefore a list of output feature numbers corresponding to
+        each encoder should be returned by this method.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def get_encoder_names(cls) -> List[str]:
+        """
+        Get the name string of encoders which will be used to initialize
+        flexible unet.
+        The reason that this function should return a list is that a
+        series of encoders can be implemented by one encoder class
+        given different initialization parameters. And a name string is
+        the key to each encoder in flexible unet backbone registry.
+        Therefore this method should return every encoder name that needs
+        to be registed in flexible unet.
+        """
+        raise NotImplementedError

monai/networks/nets/__init__.py

@@ -38,10 +38,11 @@
     EfficientNet,
     EfficientNetBN,
     EfficientNetBNFeatures,
+    EfficientNetEncoder,
     drop_connect,
     get_efficientnet_image_size,
 )
-from .flexible_unet import FlexibleUNet
+from .flexible_unet import FLEXUNET_BACKBONE, FlexibleUNet, FlexUNet, FlexUNetEncoderRegister
 from .fullyconnectednet import FullyConnectedNet, VarFullyConnectedNet
 from .generator import Generator
 from .highresnet import HighResBlock, HighResNet
@@ -50,7 +51,18 @@
 from .netadapter import NetAdapter
 from .regressor import Regressor
 from .regunet import GlobalNet, LocalNet, RegUNet
-from .resnet import ResNet, resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200
+from .resnet import (
+    ResNet,
+    ResNetBlock,
+    ResNetBottleneck,
+    resnet10,
+    resnet18,
+    resnet34,
+    resnet50,
+    resnet101,
+    resnet152,
+    resnet200,
+)
 from .segresnet import SegResNet, SegResNetVAE
 from .segresnet_ds import SegResNetDS
 from .senet import (

monai/networks/nets/efficientnet.py

@@ -13,12 +13,13 @@
 import operator
 import re
 from functools import reduce
-from typing import List, NamedTuple, Optional, Tuple, Type, Union
+from typing import Dict, List, NamedTuple, Optional, Tuple, Type, Union
 
 import torch
 from torch import nn
 from torch.utils import model_zoo
 
+from monai.networks.blocks import BaseEncoder
 from monai.networks.layers.factories import Act, Conv, Pad, Pool
 from monai.networks.layers.utils import get_norm_layer
 from monai.utils.module import look_up_option
@@ -30,6 +31,7 @@
     "drop_connect",
     "EfficientNetBNFeatures",
     "BlockArgs",
+    "EfficientNetEncoder",
 ]
 
 efficientnet_params = {
@@ -528,11 +530,8 @@ def __init__(
 
         # check if model_name is valid model
         if model_name not in efficientnet_params.keys():
-            raise ValueError(
-                "invalid model_name {} found, must be one of {} ".format(
-                    model_name, ", ".join(efficientnet_params.keys())
-                )
-            )
+            model_name_string = ", ".join(efficientnet_params.keys())
+            raise ValueError(f"invalid model_name {model_name} found, must be one of {model_name_string} ")
 
         # get network parameters
         weight_coeff, depth_coeff, image_size, dropout_rate, dropconnect_rate = efficientnet_params[model_name]
@@ -588,11 +587,8 @@ def __init__(
 
         # check if model_name is valid model
         if model_name not in efficientnet_params.keys():
-            raise ValueError(
-                "invalid model_name {} found, must be one of {} ".format(
-                    model_name, ", ".join(efficientnet_params.keys())
-                )
-            )
+            model_name_string = ", ".join(efficientnet_params.keys())
+            raise ValueError(f"invalid model_name {model_name} found, must be one of {model_name_string} ")
 
         # get network parameters
         weight_coeff, depth_coeff, image_size, dropout_rate, dropconnect_rate = efficientnet_params[model_name]
@@ -638,6 +634,80 @@ def forward(self, inputs: torch.Tensor):
         return features
 
 
+class EfficientNetEncoder(EfficientNetBNFeatures, BaseEncoder):
+    """
+    Wrap the original efficientnet to an encoder for flexible-unet.
+    """
+
+    backbone_names = [
+        "efficientnet-b0",
+        "efficientnet-b1",
+        "efficientnet-b2",
+        "efficientnet-b3",
+        "efficientnet-b4",
+        "efficientnet-b5",
+        "efficientnet-b6",
+        "efficientnet-b7",
+        "efficientnet-b8",
+        "efficientnet-l2",
+    ]
+
+    @classmethod
+    def get_encoder_parameters(cls) -> List[Dict]:
+        """
+        Get the initialization parameter for efficientnet backbones.
+        """
+        parameter_list = []
+        for backbone_name in cls.backbone_names:
+            parameter_list.append(
+                {
+                    "model_name": backbone_name,
+                    "pretrained": True,
+                    "progress": True,
+                    "spatial_dims": 2,
+                    "in_channels": 3,
+                    "num_classes": 1000,
+                    "norm": ("batch", {"eps": 1e-3, "momentum": 0.01}),
+                    "adv_prop": "ap" in backbone_name,
+                }
+            )
+        return parameter_list
+
+    @classmethod
+    def num_channels_per_output(cls) -> List[Tuple[int, ...]]:
+        """
+        Get number of efficientnet backbone output feature maps' channel.
+        """
+        return [
+            (16, 24, 40, 112, 320),
+            (16, 24, 40, 112, 320),
+            (16, 24, 48, 120, 352),
+            (24, 32, 48, 136, 384),
+            (24, 32, 56, 160, 448),
+            (24, 40, 64, 176, 512),
+            (32, 40, 72, 200, 576),
+            (32, 48, 80, 224, 640),
+            (32, 56, 88, 248, 704),
+            (72, 104, 176, 480, 1376),
+        ]
+
+    @classmethod
+    def num_outputs(cls) -> List[int]:
+        """
+        Get number of efficientnet backbone output feature maps.
+        Since every backbone contains the same 5 output feature maps,
+        the number list should be `[5] * 10`.
+        """
+        return [5] * 10
+
+    @classmethod
+    def get_encoder_names(cls) -> List[str]:
+        """
+        Get names of efficient backbone.
+        """
+        return cls.backbone_names
+
+
 def get_efficientnet_image_size(model_name: str) -> int:
     """
     Get the input image size for a given efficientnet model.
@@ -651,9 +721,8 @@ def get_efficientnet_image_size(model_name: str) -> int:
     """
     # check if model_name is valid model
     if model_name not in efficientnet_params.keys():
-        raise ValueError(
-            "invalid model_name {} found, must be one of {} ".format(model_name, ", ".join(efficientnet_params.keys()))
-        )
+        model_name_string = ", ".join(efficientnet_params.keys())
+        raise ValueError(f"invalid model_name {model_name} found, must be one of {model_name_string} ")
 
     # return input image size (all dims equal so only need to return for one dim)
     _, _, res, _, _ = efficientnet_params[model_name]
@@ -927,15 +996,10 @@ def to_string(self):
             A string notation of BlockArgs object arguments.
                 Example: "r1_k3_s11_e1_i32_o16_se0.25_noskip".
         """
-        string = "r{}_k{}_s{}{}_e{}_i{}_o{}_se{}".format(
-            self.num_repeat,
-            self.kernel_size,
-            self.stride,
-            self.stride,
-            self.expand_ratio,
-            self.input_filters,
-            self.output_filters,
-            self.se_ratio,
+        string = (
+            f"r{self.num_repeat}_k{self.kernel_size}_s{self.stride}{self.stride}"
+            f"_e{self.expand_ratio}_i{self.input_filters}_o{self.output_filters}"
+            f"_se{self.se_ratio}"
         )
 
         if not self.id_skip: