Make the flexible unet more extensible to different backbones (Project-MONAI#5038)

Signed-off-by: binliu <binliu@nvidia.com> Fixes Project-MONAI#5011 . ### Description Currently, the FlexibleUNet structure only supports the efficient-net series as the backbone and it is hard to extend to other network structures. A more extensible and convenient way to add more benchmark or user-defined backbones will make it more flexible. I plan to make it by doing steps below 1. A base backbone/encoder class that defines interfaces like the number of output feature maps, a list of output feature map channels, string names of backbones and so on. 2. A register that can dynamically register backbones to a dict that will be used during the flexible unet initialization. ### Types of changes <!--- Put an `x` in all the boxes that apply, and remove the not applicable items --> - [x] Non-breaking change (fix or new feature that would not break existing functionality). - [ ] Breaking change (fix or new feature that would cause existing functionality to change). - [ ] New tests added to cover the changes. - [ ] Integration tests passed locally by running `./runtests.sh -f -u --net --coverage`. - [ ] Quick tests passed locally by running `./runtests.sh --quick --unittests --disttests`. - [ ] In-line docstrings updated. - [ ] Documentation updated, tested `make html` command in the `docs/` folder. Signed-off-by: binliu <binliu@nvidia.com>

Author: binliunls

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: