def __init__(
self,
spatial_dims: int,
in_channels: int,
norm: Union[Tuple, str],
kernel_size: int = 3,
) -> None:
"""
Args:
spatial_dims: number of spatial dimensions, could be 1, 2 or 3.
in_channels: number of input channels.
norm: feature normalization type and arguments.
kernel_size: convolution kernel size, the value should be an odd number. Defaults to 3.
"""
super().__init__()
if kernel_size % 2 != 1:
raise AssertionError("kernel_size should be an odd number.")
self.norm1 = get_norm_layer(name=norm,
spatial_dims=spatial_dims,
channels=in_channels)
self.norm2 = get_norm_layer(name=norm,
spatial_dims=spatial_dims,
channels=in_channels)
self.relu = Act[Act.RELU](inplace=True)
self.conv1 = get_conv_layer(spatial_dims,
in_channels=in_channels,
out_channels=in_channels)
self.conv2 = get_conv_layer(spatial_dims,
in_channels=in_channels,
out_channels=in_channels)
def __init__(
self,
num_features: int,
in_channels: int,
out_channels: int,
dropout_prob: float = 0.0,
act: Union[str, tuple] = ("relu", {
"inplace": True
}),
norm: Union[str, tuple] = "batch",
kernel_size: int = 3,
) -> None:
"""
Args:
spatial_dims: number of spatial dimensions of the input image.
num_features: number of internal channels used for the layer
in_channels: number of the input channels.
out_channels: number of the output channels.
dropout_prob: dropout rate after each dense layer.
act: activation type and arguments. Defaults to relu.
norm: feature normalization type and arguments. Defaults to batch norm.
kernel_size: size of the kernel for >1 convolutions (dependent on mode)
"""
super().__init__()
conv_type: Callable = Conv[Conv.CONV, 2]
dropout_type: Callable = Dropout[Dropout.DROPOUT, 2]
self.layers = nn.Sequential()
self.layers.add_module(
"preact_bna/bn",
get_norm_layer(name=norm, spatial_dims=2, channels=in_channels))
self.layers.add_module("preact_bna/relu", get_act_layer(name=act))
self.layers.add_module(
"conv1",
conv_type(in_channels, num_features, kernel_size=1, bias=False))
self.layers.add_module(
"conv1/norm",
get_norm_layer(name=norm, spatial_dims=2, channels=num_features))
self.layers.add_module("conv1/relu2", get_act_layer(name=act))
self.layers.add_module(
"conv2",
conv_type(num_features,
out_channels,
kernel_size=kernel_size,
padding=0,
groups=4,
bias=False))
if dropout_prob > 0:
self.layers.add_module("dropout", dropout_type(dropout_prob))
def __init__(
self,
spatial_dims: int,
in_channels: int,
out_channels: int,
kernel_size: Union[Sequence[int], int],
stride: Union[Sequence[int], int],
norm_name: Union[tuple, str],
act_name: Union[tuple, str] = ("leakyrelu", {
"inplace": True,
"negative_slope": 0.01
}),
dropout: Optional[Union[tuple, str, float]] = None,
):
super().__init__()
self.conv1 = get_conv_layer(
spatial_dims,
in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
dropout=dropout,
conv_only=True,
)
self.conv2 = get_conv_layer(spatial_dims,
out_channels,
out_channels,
kernel_size=kernel_size,
stride=1,
dropout=dropout,
conv_only=True)
self.conv3 = get_conv_layer(spatial_dims,
in_channels,
out_channels,
kernel_size=1,
stride=stride,
dropout=dropout,
conv_only=True)
self.lrelu = get_act_layer(name=act_name)
self.norm1 = get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=out_channels)
self.norm2 = get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=out_channels)
self.norm3 = get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=out_channels)
self.downsample = in_channels != out_channels
stride_np = np.atleast_1d(stride)
if not np.all(stride_np == 1):
self.downsample = True
def __init__(
self,
spatial_dims: int,
in_channels: int,
out_channels: int,
kernel_size: Union[Sequence[int], int],
stride: Union[Sequence[int], int],
norm_name: Union[Tuple, str],
):
super(UnetResBlock, self).__init__()
self.conv1 = get_conv_layer(
spatial_dims,
in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
conv_only=True,
)
self.conv2 = get_conv_layer(
spatial_dims,
out_channels,
out_channels,
kernel_size=kernel_size,
stride=1,
conv_only=True,
)
self.conv3 = get_conv_layer(
spatial_dims,
in_channels,
out_channels,
kernel_size=1,
stride=stride,
conv_only=True,
)
self.lrelu = get_act_layer(("leakyrelu", {
"inplace": True,
"negative_slope": 0.01
}))
self.norm1 = get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=out_channels)
self.norm2 = get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=out_channels)
self.norm3 = get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=out_channels)
self.downsample = in_channels != out_channels
stride_np = np.atleast_1d(stride)
if not np.all(stride_np == 1):
self.downsample = True
def __init__(
self,
spatial_dims: int,
in_channels: int,
out_channels: int,
act: Union[str, tuple] = ("relu", {"inplace": True}),
norm: Union[str, tuple] = "batch",
) -> None:
"""
Args:
spatial_dims: number of spatial dimensions of the input image.
in_channels: number of the input channel.
out_channels: number of the output classes.
act: activation type and arguments. Defaults to relu.
norm: feature normalization type and arguments. Defaults to batch norm.
"""
super().__init__()
conv_type: Callable = Conv[Conv.CONV, spatial_dims]
pool_type: Callable = Pool[Pool.AVG, spatial_dims]
self.add_module("norm", get_norm_layer(name=norm, spatial_dims=spatial_dims, channels=in_channels))
self.add_module("relu", get_act_layer(name=act))
self.add_module("conv", conv_type(in_channels, out_channels, kernel_size=1, bias=False))
self.add_module("pool", pool_type(kernel_size=2, stride=2))
def __init__(
self,
ordering: str = "NDA",
in_channels: Optional[int] = None,
act: Optional[Union[Tuple, str]] = "RELU",
norm: Optional[Union[Tuple, str]] = None,
norm_dim: Optional[int] = None,
dropout: Optional[Union[Tuple, str, float]] = None,
dropout_dim: Optional[int] = None,
) -> None:
super().__init__()
op_dict = {"A": None, "D": None, "N": None}
# define the normalization type and the arguments to the constructor
if norm is not None:
if norm_dim is None and dropout_dim is None:
raise ValueError("norm_dim or dropout_dim needs to be specified.")
op_dict["N"] = get_norm_layer(name=norm, spatial_dims=norm_dim or dropout_dim, channels=in_channels)
# define the activation type and the arguments to the constructor
if act is not None:
op_dict["A"] = get_act_layer(act)
if dropout is not None:
if norm_dim is None and dropout_dim is None:
raise ValueError("norm_dim or dropout_dim needs to be specified.")
op_dict["D"] = get_dropout_layer(name=dropout, dropout_dim=dropout_dim or norm_dim)
for item in ordering.upper():
if item not in op_dict:
raise ValueError(f"ordering must be a string of {op_dict}, got {item} in it.")
if op_dict[item] is not None:
self.add_module(item, op_dict[item])
def __init__(
self,
spatial_dims: int,
in_channels: int,
out_channels: int,
kernel_size: Union[Sequence[int], int],
stride: Union[Sequence[int], int],
norm_name: Union[Tuple, str],
act_name: Union[Tuple, str] = ("leakyrelu", {
"inplace": True,
"negative_slope": 0.01
}),
dropout: Optional[Union[Tuple, str, float]] = None,
):
super().__init__()
self.conv1 = get_conv_layer(
spatial_dims,
in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
dropout=dropout,
act=None,
norm=None,
conv_only=False,
)
self.conv2 = get_conv_layer(
spatial_dims,
out_channels,
out_channels,
kernel_size=kernel_size,
stride=1,
dropout=dropout,
act=None,
norm=None,
conv_only=False,
)
self.lrelu = get_act_layer(name=act_name)
self.norm1 = get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=out_channels)
self.norm2 = get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=out_channels)
def _make_final_conv(self, out_channels: int):
return nn.Sequential(
get_norm_layer(name=self.norm,
spatial_dims=self.spatial_dims,
channels=self.init_filters),
self.act,
get_conv_layer(self.spatial_dims,
self.init_filters,
out_channels,
kernel_size=1,
bias=True),
)
def __init__(
self,
spatial_dims: int,
in_channels: int,
growth_rate: int,
bn_size: int,
dropout_prob: float,
act: Union[str, tuple] = ("relu", {"inplace": True}),
norm: Union[str, tuple] = "batch",
) -> None:
"""
Args:
spatial_dims: number of spatial dimensions of the input image.
in_channels: number of the input channel.
growth_rate: how many filters to add each layer (k in paper).
bn_size: multiplicative factor for number of bottle neck layers.
(i.e. bn_size * k features in the bottleneck layer)
dropout_prob: dropout rate after each dense layer.
act: activation type and arguments. Defaults to relu.
norm: feature normalization type and arguments. Defaults to batch norm.
"""
super().__init__()
out_channels = bn_size * growth_rate
conv_type: Callable = Conv[Conv.CONV, spatial_dims]
dropout_type: Callable = Dropout[Dropout.DROPOUT, spatial_dims]
self.layers = nn.Sequential()
self.layers.add_module("norm1", get_norm_layer(name=norm, spatial_dims=spatial_dims, channels=in_channels))
self.layers.add_module("relu1", get_act_layer(name=act))
self.layers.add_module("conv1", conv_type(in_channels, out_channels, kernel_size=1, bias=False))
self.layers.add_module("norm2", get_norm_layer(name=norm, spatial_dims=spatial_dims, channels=out_channels))
self.layers.add_module("relu2", get_act_layer(name=act))
self.layers.add_module("conv2", conv_type(out_channels, growth_rate, kernel_size=3, padding=1, bias=False))
if dropout_prob > 0:
self.layers.add_module("dropout", dropout_type(dropout_prob))
def __init__(
self,
in_channel: int,
out_channel: int,
spatial_dims: int = 3,
act_name: Union[Tuple, str] = "RELU",
norm_name: Union[Tuple, str] = ("INSTANCE", {
"affine": True
}),
):
"""
Args:
in_channel: number of input channels
out_channel: number of output channels.
spatial_dims: number of spatial dimensions.
act_name: activation layer type and arguments.
norm_name: feature normalization type and arguments.
"""
super().__init__()
self._in_channel = in_channel
self._out_channel = out_channel
self._spatial_dims = spatial_dims
if self._spatial_dims not in (2, 3):
raise ValueError("spatial_dims must be 2 or 3.")
conv_type = Conv[Conv.CONV, self._spatial_dims]
self.act = get_act_layer(name=act_name)
self.conv_1 = conv_type(
in_channels=self._in_channel,
out_channels=self._out_channel // 2,
kernel_size=1,
stride=2,
padding=0,
groups=1,
bias=False,
dilation=1,
)
self.conv_2 = conv_type(
in_channels=self._in_channel,
out_channels=self._out_channel - self._out_channel // 2,
kernel_size=1,
stride=2,
padding=0,
groups=1,
bias=False,
dilation=1,
)
self.norm = get_norm_layer(name=norm_name,
spatial_dims=self._spatial_dims,
channels=self._out_channel)
def _prepare_vae_modules(self):
zoom = 2**(len(self.blocks_down) - 1)
v_filters = self.init_filters * zoom
total_elements = int(self.smallest_filters * np.prod(self.fc_insize))
self.vae_down = nn.Sequential(
get_norm_layer(name=self.norm,
spatial_dims=self.spatial_dims,
channels=v_filters),
self.act,
get_conv_layer(self.spatial_dims,
v_filters,
self.smallest_filters,
stride=2,
bias=True),
get_norm_layer(name=self.norm,
spatial_dims=self.spatial_dims,
channels=self.smallest_filters),
self.act,
)
self.vae_fc1 = nn.Linear(total_elements, self.vae_nz)
self.vae_fc2 = nn.Linear(total_elements, self.vae_nz)
self.vae_fc3 = nn.Linear(self.vae_nz, total_elements)
self.vae_fc_up_sample = nn.Sequential(
get_conv_layer(self.spatial_dims,
self.smallest_filters,
v_filters,
kernel_size=1),
get_upsample_layer(self.spatial_dims,
v_filters,
upsample_mode=self.upsample_mode),
get_norm_layer(name=self.norm,
spatial_dims=self.spatial_dims,
channels=v_filters),
self.act,
)
def __init__(
self,
in_channel: int,
out_channel: int,
spatial_dims: int = 3,
act_name: Union[Tuple, str] = "RELU",
norm_name: Union[Tuple, str] = ("INSTANCE", {
"affine": True
}),
):
"""
Args:
in_channel: number of input channels
out_channel: number of output channels
spatial_dims: number of spatial dimensions
act_name: activation layer type and arguments.
norm_name: feature normalization type and arguments.
"""
super().__init__()
self._in_channel = in_channel
self._out_channel = out_channel
self._spatial_dims = spatial_dims
if self._spatial_dims not in (2, 3):
raise ValueError("spatial_dims must be 2 or 3.")
conv_type = Conv[Conv.CONV, self._spatial_dims]
mode = "trilinear" if self._spatial_dims == 3 else "bilinear"
self.add_module(
"up",
torch.nn.Upsample(scale_factor=2, mode=mode, align_corners=True))
self.add_module("acti", get_act_layer(name=act_name))
self.add_module(
"conv",
conv_type(
in_channels=self._in_channel,
out_channels=self._out_channel,
kernel_size=1,
stride=1,
padding=0,
groups=1,
bias=False,
dilation=1,
),
)
self.add_module(
"norm",
get_norm_layer(name=norm_name,
spatial_dims=self._spatial_dims,
channels=self._out_channel))
def __init__(
self,
spatial_dims: int,
in_channels: int,
out_channels: int,
kernel_size: Union[Sequence[int], int],
stride: Union[Sequence[int], int],
norm_name: Union[Tuple, str],
):
super(UnetBasicBlock, self).__init__()
self.conv1 = get_conv_layer(
spatial_dims,
in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
conv_only=True,
)
self.conv2 = get_conv_layer(
spatial_dims,
out_channels,
out_channels,
kernel_size=kernel_size,
stride=1,
conv_only=True,
)
self.lrelu = get_act_layer(("leakyrelu", {
"inplace": True,
"negative_slope": 0.01
}))
self.norm1 = get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=out_channels)
self.norm2 = get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=out_channels)
def __init__(
self,
in_channel: int,
out_channel: int,
kernel_size: int = 3,
padding: int = 1,
spatial_dims: int = 3,
act_name: Union[Tuple, str] = "RELU",
norm_name: Union[Tuple, str] = ("INSTANCE", {
"affine": True
}),
):
"""
Args:
in_channel: number of input channels.
out_channel: number of output channels.
kernel_size: kernel size of the convolution.
padding: padding size of the convolution.
spatial_dims: number of spatial dimensions.
act_name: activation layer type and arguments.
norm_name: feature normalization type and arguments.
"""
super().__init__()
self._in_channel = in_channel
self._out_channel = out_channel
self._spatial_dims = spatial_dims
conv_type = Conv[Conv.CONV, self._spatial_dims]
self.add_module("acti", get_act_layer(name=act_name))
self.add_module(
"conv",
conv_type(
in_channels=self._in_channel,
out_channels=self._out_channel,
kernel_size=kernel_size,
stride=1,
padding=padding,
groups=1,
bias=False,
dilation=1,
),
)
self.add_module(
"norm",
get_norm_layer(name=norm_name,
spatial_dims=self._spatial_dims,
channels=self._out_channel))
def __init__(self,
in_channels: int,
act: Union[str, tuple] = ("relu", {
"inplace": True
}),
norm: Union[str, tuple] = "batch") -> None:
"""
Args:
in_channels: number of the input channel.
act: activation type and arguments. Defaults to relu.
norm: feature normalization type and arguments. Defaults to batch norm.
"""
super().__init__()
self.add_module(
"norm",
get_norm_layer(name=norm, spatial_dims=2, channels=in_channels))
self.add_module("relu", get_act_layer(name=act))
def __init__(
self,
dints_space,
in_channels: int,
num_classes: int,
act_name: Union[Tuple, str] = "RELU",
norm_name: Union[Tuple, str] = ("INSTANCE", {
"affine": True
}),
spatial_dims: int = 3,
use_downsample: bool = True,
node_a=None,
):
super().__init__()
self.dints_space = dints_space
self.filter_nums = dints_space.filter_nums
self.num_blocks = dints_space.num_blocks
self.num_depths = dints_space.num_depths
if spatial_dims not in (2, 3):
raise NotImplementedError(
f"Spatial dimensions {spatial_dims} is not supported.")
self._spatial_dims = spatial_dims
if node_a is None:
self.node_a = torch.ones((self.num_blocks + 1, self.num_depths))
else:
self.node_a = node_a
# define stem operations for every block
conv_type = Conv[Conv.CONV, spatial_dims]
self.stem_down = nn.ModuleDict()
self.stem_up = nn.ModuleDict()
self.stem_finals = nn.Sequential(
ActiConvNormBlock(
self.filter_nums[0],
self.filter_nums[0],
act_name=act_name,
norm_name=norm_name,
spatial_dims=spatial_dims,
),
conv_type(
in_channels=self.filter_nums[0],
out_channels=num_classes,
kernel_size=1,
stride=1,
padding=0,
groups=1,
bias=True,
dilation=1,
),
)
mode = "trilinear" if self._spatial_dims == 3 else "bilinear"
for res_idx in range(self.num_depths):
# define downsample stems before DiNTS search
if use_downsample:
self.stem_down[str(res_idx)] = StemTS(
nn.Upsample(scale_factor=1 / (2**res_idx),
mode=mode,
align_corners=True),
conv_type(
in_channels=in_channels,
out_channels=self.filter_nums[res_idx],
kernel_size=3,
stride=1,
padding=1,
groups=1,
bias=False,
dilation=1,
),
get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=self.filter_nums[res_idx]),
get_act_layer(name=act_name),
conv_type(
in_channels=self.filter_nums[res_idx],
out_channels=self.filter_nums[res_idx + 1],
kernel_size=3,
stride=2,
padding=1,
groups=1,
bias=False,
dilation=1,
),
get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=self.filter_nums[res_idx + 1]),
)
self.stem_up[str(res_idx)] = StemTS(
get_act_layer(name=act_name),
conv_type(
in_channels=self.filter_nums[res_idx + 1],
out_channels=self.filter_nums[res_idx],
kernel_size=3,
stride=1,
padding=1,
groups=1,
bias=False,
dilation=1,
),
get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=self.filter_nums[res_idx]),
nn.Upsample(scale_factor=2, mode=mode, align_corners=True),
)
else:
self.stem_down[str(res_idx)] = StemTS(
nn.Upsample(scale_factor=1 / (2**res_idx),
mode=mode,
align_corners=True),
conv_type(
in_channels=in_channels,
out_channels=self.filter_nums[res_idx],
kernel_size=3,
stride=1,
padding=1,
groups=1,
bias=False,
dilation=1,
),
get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=self.filter_nums[res_idx]),
)
self.stem_up[str(res_idx)] = StemTS(
get_act_layer(name=act_name),
conv_type(
in_channels=self.filter_nums[res_idx],
out_channels=self.filter_nums[max(res_idx - 1, 0)],
kernel_size=3,
stride=1,
padding=1,
groups=1,
bias=False,
dilation=1,
),
get_norm_layer(name=norm_name,
spatial_dims=spatial_dims,
channels=self.filter_nums[max(
res_idx - 1, 0)]),
nn.Upsample(scale_factor=2**(res_idx != 0),
mode=mode,
align_corners=True),
)
def __init__(
self,
blocks_args_str: List[str],
spatial_dims: int = 2,
in_channels: int = 3,
num_classes: int = 1000,
width_coefficient: float = 1.0,
depth_coefficient: float = 1.0,
dropout_rate: float = 0.2,
image_size: int = 224,
norm: Union[str, tuple] = ("batch", {
"eps": 1e-3,
"momentum": 0.01
}),
drop_connect_rate: float = 0.2,
depth_divisor: int = 8,
) -> None:
"""
EfficientNet based on `Rethinking Model Scaling for Convolutional Neural Networks <https://arxiv.org/pdf/1905.11946.pdf>`_.
Adapted from `EfficientNet-PyTorch <https://github.com/lukemelas/EfficientNet-PyTorch>`_.
Args:
blocks_args_str: block definitions.
spatial_dims: number of spatial dimensions.
in_channels: number of input channels.
num_classes: number of output classes.
width_coefficient: width multiplier coefficient (w in paper).
depth_coefficient: depth multiplier coefficient (d in paper).
dropout_rate: dropout rate for dropout layers.
image_size: input image resolution.
norm: feature normalization type and arguments.
drop_connect_rate: dropconnect rate for drop connection (individual weights) layers.
depth_divisor: depth divisor for channel rounding.
"""
super().__init__()
if spatial_dims not in (1, 2, 3):
raise ValueError("spatial_dims can only be 1, 2 or 3.")
# select the type of N-Dimensional layers to use
# these are based on spatial dims and selected from MONAI factories
conv_type: Type[Union[nn.Conv1d, nn.Conv2d,
nn.Conv3d]] = Conv["conv", spatial_dims]
adaptivepool_type: Type[
Union[nn.AdaptiveAvgPool1d, nn.AdaptiveAvgPool2d,
nn.AdaptiveAvgPool3d]] = Pool["adaptiveavg", spatial_dims]
# decode blocks args into arguments for MBConvBlock
blocks_args = [BlockArgs.from_string(s) for s in blocks_args_str]
# checks for successful decoding of blocks_args_str
if not isinstance(blocks_args, list):
raise ValueError("blocks_args must be a list")
if blocks_args == []:
raise ValueError("block_args must be non-empty")
self._blocks_args = blocks_args
self.num_classes = num_classes
self.in_channels = in_channels
self.drop_connect_rate = drop_connect_rate
# expand input image dimensions to list
current_image_size = [image_size] * spatial_dims
# Stem
stride = 2
out_channels = _round_filters(
32, width_coefficient, depth_divisor) # number of output channels
self._conv_stem = conv_type(self.in_channels,
out_channels,
kernel_size=3,
stride=stride,
bias=False)
self._conv_stem_padding = _make_same_padder(self._conv_stem,
current_image_size)
self._bn0 = get_norm_layer(name=norm,
spatial_dims=spatial_dims,
channels=out_channels)
current_image_size = _calculate_output_image_size(
current_image_size, stride)
# build MBConv blocks
num_blocks = 0
self._blocks = nn.Sequential()
self.extract_stacks = []
# update baseline blocks to input/output filters and number of repeats based on width and depth multipliers.
for idx, block_args in enumerate(self._blocks_args):
block_args = block_args._replace(
input_filters=_round_filters(block_args.input_filters,
width_coefficient, depth_divisor),
output_filters=_round_filters(block_args.output_filters,
width_coefficient,
depth_divisor),
num_repeat=_round_repeats(block_args.num_repeat,
depth_coefficient),
)
self._blocks_args[idx] = block_args
# calculate the total number of blocks - needed for drop_connect estimation
num_blocks += block_args.num_repeat
if block_args.stride > 1:
self.extract_stacks.append(idx)
self.extract_stacks.append(len(self._blocks_args))
# create and add MBConvBlocks to self._blocks
idx = 0 # block index counter
for stack_idx, block_args in enumerate(self._blocks_args):
blk_drop_connect_rate = self.drop_connect_rate
# scale drop connect_rate
if blk_drop_connect_rate:
blk_drop_connect_rate *= float(idx) / num_blocks
sub_stack = nn.Sequential()
# the first block needs to take care of stride and filter size increase.
sub_stack.add_module(
str(idx),
MBConvBlock(
spatial_dims=spatial_dims,
in_channels=block_args.input_filters,
out_channels=block_args.output_filters,
kernel_size=block_args.kernel_size,
stride=block_args.stride,
image_size=current_image_size,
expand_ratio=block_args.expand_ratio,
se_ratio=block_args.se_ratio,
id_skip=block_args.id_skip,
norm=norm,
drop_connect_rate=blk_drop_connect_rate,
),
)
idx += 1 # increment blocks index counter
current_image_size = _calculate_output_image_size(
current_image_size, block_args.stride)
if block_args.num_repeat > 1: # modify block_args to keep same output size
block_args = block_args._replace(
input_filters=block_args.output_filters, stride=1)
# add remaining block repeated num_repeat times
for _ in range(block_args.num_repeat - 1):
blk_drop_connect_rate = self.drop_connect_rate
# scale drop connect_rate
if blk_drop_connect_rate:
blk_drop_connect_rate *= float(idx) / num_blocks
# add blocks
sub_stack.add_module(
str(idx),
MBConvBlock(
spatial_dims=spatial_dims,
in_channels=block_args.input_filters,
out_channels=block_args.output_filters,
kernel_size=block_args.kernel_size,
stride=block_args.stride,
image_size=current_image_size,
expand_ratio=block_args.expand_ratio,
se_ratio=block_args.se_ratio,
id_skip=block_args.id_skip,
norm=norm,
drop_connect_rate=blk_drop_connect_rate,
),
)
idx += 1 # increment blocks index counter
self._blocks.add_module(str(stack_idx), sub_stack)
# sanity check to see if len(self._blocks) equal expected num_blocks
if idx != num_blocks:
raise ValueError("total number of blocks created != num_blocks")
# Head
head_in_channels = block_args.output_filters
out_channels = _round_filters(1280, width_coefficient, depth_divisor)
self._conv_head = conv_type(head_in_channels,
out_channels,
kernel_size=1,
bias=False)
self._conv_head_padding = _make_same_padder(self._conv_head,
current_image_size)
self._bn1 = get_norm_layer(name=norm,
spatial_dims=spatial_dims,
channels=out_channels)
# final linear layer
self._avg_pooling = adaptivepool_type(1)
self._dropout = nn.Dropout(dropout_rate)
self._fc = nn.Linear(out_channels, self.num_classes)
# swish activation to use - using memory efficient swish by default
# can be switched to normal swish using self.set_swish() function call
self._swish = Act["memswish"]()
# initialize weights using Tensorflow's init method from official impl.
self._initialize_weights()
def __init__(
self,
spatial_dims: int,
in_channels: int,
out_channels: int,
init_features: int = 64,
growth_rate: int = 32,
block_config: Sequence[int] = (6, 12, 24, 16),
bn_size: int = 4,
act: Union[str, tuple] = ("relu", {"inplace": True}),
norm: Union[str, tuple] = "batch",
dropout_prob: float = 0.0,
) -> None:
super().__init__()
conv_type: Type[Union[nn.Conv1d, nn.Conv2d, nn.Conv3d]] = Conv[Conv.CONV, spatial_dims]
pool_type: Type[Union[nn.MaxPool1d, nn.MaxPool2d, nn.MaxPool3d]] = Pool[Pool.MAX, spatial_dims]
avg_pool_type: Type[Union[nn.AdaptiveAvgPool1d, nn.AdaptiveAvgPool2d, nn.AdaptiveAvgPool3d]] = Pool[
Pool.ADAPTIVEAVG, spatial_dims
]
self.features = nn.Sequential(
OrderedDict(
[
("conv0", conv_type(in_channels, init_features, kernel_size=7, stride=2, padding=3, bias=False)),
("norm0", get_norm_layer(name=norm, spatial_dims=spatial_dims, channels=init_features)),
("relu0", get_act_layer(name=act)),
("pool0", pool_type(kernel_size=3, stride=2, padding=1)),
]
)
)
in_channels = init_features
for i, num_layers in enumerate(block_config):
block = _DenseBlock(
spatial_dims=spatial_dims,
layers=num_layers,
in_channels=in_channels,
bn_size=bn_size,
growth_rate=growth_rate,
dropout_prob=dropout_prob,
act=act,
norm=norm,
)
self.features.add_module(f"denseblock{i + 1}", block)
in_channels += num_layers * growth_rate
if i == len(block_config) - 1:
self.features.add_module(
"norm5", get_norm_layer(name=norm, spatial_dims=spatial_dims, channels=in_channels)
)
else:
_out_channels = in_channels // 2
trans = _Transition(
spatial_dims, in_channels=in_channels, out_channels=_out_channels, act=act, norm=norm
)
self.features.add_module(f"transition{i + 1}", trans)
in_channels = _out_channels
# pooling and classification
self.class_layers = nn.Sequential(
OrderedDict(
[
("relu", get_act_layer(name=act)),
("pool", avg_pool_type(1)),
("flatten", nn.Flatten(1)),
("out", nn.Linear(in_channels, out_channels)),
]
)
)
for m in self.modules():
if isinstance(m, conv_type):
nn.init.kaiming_normal_(torch.as_tensor(m.weight))
elif isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
nn.init.constant_(torch.as_tensor(m.weight), 1)
nn.init.constant_(torch.as_tensor(m.bias), 0)
elif isinstance(m, nn.Linear):
nn.init.constant_(torch.as_tensor(m.bias), 0)
def __init__(
self,
spatial_dims: int,
in_channels: int,
out_channels: int,
kernel_size: int,
stride: int,
image_size: List[int],
expand_ratio: int,
se_ratio: Optional[float],
id_skip: Optional[bool] = True,
norm: Union[str, tuple] = ("batch", {
"eps": 1e-3,
"momentum": 0.01
}),
drop_connect_rate: Optional[float] = 0.2,
) -> None:
"""
Mobile Inverted Residual Bottleneck Block.
Args:
spatial_dims: number of spatial dimensions.
in_channels: number of input channels.
out_classes: number of output channels.
kernel_size: size of the kernel for conv ops.
stride: stride to use for conv ops.
image_size: input image resolution.
expand_ratio: expansion ratio for inverted bottleneck.
se_ratio: squeeze-excitation ratio for se layers.
id_skip: whether to use skip connection.
norm: feature normalization type and arguments. Defaults to batch norm.
drop_connect_rate: dropconnect rate for drop connection (individual weights) layers.
References:
[1] https://arxiv.org/abs/1704.04861 (MobileNet v1)
[2] https://arxiv.org/abs/1801.04381 (MobileNet v2)
[3] https://arxiv.org/abs/1905.02244 (MobileNet v3)
"""
super().__init__()
# select the type of N-Dimensional layers to use
# these are based on spatial dims and selected from MONAI factories
conv_type = Conv["conv", spatial_dims]
adaptivepool_type = Pool["adaptiveavg", spatial_dims]
self.in_channels = in_channels
self.out_channels = out_channels
self.id_skip = id_skip
self.stride = stride
self.expand_ratio = expand_ratio
self.drop_connect_rate = drop_connect_rate
if (se_ratio is not None) and (0.0 < se_ratio <= 1.0):
self.has_se = True
self.se_ratio = se_ratio
else:
self.has_se = False
# Expansion phase (Inverted Bottleneck)
inp = in_channels # number of input channels
oup = in_channels * expand_ratio # number of output channels
if self.expand_ratio != 1:
self._expand_conv = conv_type(in_channels=inp,
out_channels=oup,
kernel_size=1,
bias=False)
self._expand_conv_padding = _make_same_padder(
self._expand_conv, image_size)
self._bn0 = get_norm_layer(name=norm,
spatial_dims=spatial_dims,
channels=oup)
else:
# need to have the following to fix JIT error:
# "Module 'MBConvBlock' has no attribute '_expand_conv'"
# FIXME: find a better way to bypass JIT error
self._expand_conv = nn.Identity()
self._expand_conv_padding = nn.Identity()
self._bn0 = nn.Identity()
# Depthwise convolution phase
self._depthwise_conv = conv_type(
in_channels=oup,
out_channels=oup,
groups=oup, # groups makes it depthwise
kernel_size=kernel_size,
stride=self.stride,
bias=False,
)
self._depthwise_conv_padding = _make_same_padder(
self._depthwise_conv, image_size)
self._bn1 = get_norm_layer(name=norm,
spatial_dims=spatial_dims,
channels=oup)
image_size = _calculate_output_image_size(image_size, self.stride)
# Squeeze and Excitation layer, if desired
if self.has_se:
self._se_adaptpool = adaptivepool_type(1)
num_squeezed_channels = max(1, int(in_channels * self.se_ratio))
self._se_reduce = conv_type(in_channels=oup,
out_channels=num_squeezed_channels,
kernel_size=1)
self._se_reduce_padding = _make_same_padder(
self._se_reduce, [1, 1])
self._se_expand = conv_type(in_channels=num_squeezed_channels,
out_channels=oup,
kernel_size=1)
self._se_expand_padding = _make_same_padder(
self._se_expand, [1, 1])
# Pointwise convolution phase
final_oup = out_channels
self._project_conv = conv_type(in_channels=oup,
out_channels=final_oup,
kernel_size=1,
bias=False)
self._project_conv_padding = _make_same_padder(self._project_conv,
image_size)
self._bn2 = get_norm_layer(name=norm,
spatial_dims=spatial_dims,
channels=final_oup)
# swish activation to use - using memory efficient swish by default
# can be switched to normal swish using self.set_swish() function call
self._swish = Act["memswish"](inplace=True)
def __init__(
self,
mode: Mode = Mode.FAST,
in_channels: int = 3,
out_classes: int = 0,
act: Union[str, tuple] = ("relu", {
"inplace": True
}),
norm: Union[str, tuple] = "batch",
dropout_prob: float = 0.0,
) -> None:
super().__init__()
self.mode: int = self._mode_to_int(mode)
if mode not in [self.Mode.ORIGINAL, self.Mode.FAST]:
raise ValueError(
"Input size should be 270 x 270 when using Mode.ORIGINAL")
if out_classes > 128:
raise ValueError(
"Number of nuclear types classes exceeds maximum (128)")
elif out_classes == 1:
raise ValueError(
"Number of nuclear type classes should either be None or >1")
if dropout_prob > 1 or dropout_prob < 0:
raise ValueError("Dropout can only be in the range 0.0 to 1.0")
# number of filters in the first convolution layer.
_init_features: int = 64
# number of layers in each pooling block.
_block_config: Sequence[int] = (3, 4, 6, 3)
if mode == self.Mode.FAST:
_ksize = 3
_pad = 3
else:
_ksize = 5
_pad = 0
conv_type: Type[nn.Conv2d] = Conv[Conv.CONV, 2]
self.input_features = nn.Sequential(
OrderedDict([
(
"conv0",
conv_type(in_channels,
_init_features,
kernel_size=7,
stride=1,
padding=_pad,
bias=False),
),
("norm0",
get_norm_layer(name=norm,
spatial_dims=2,
channels=_init_features)),
("relu0", get_act_layer(name=act)),
]))
_in_channels = _init_features
_out_channels = 256
_num_features = _init_features
self.res_blocks = nn.Sequential()
for i, num_layers in enumerate(_block_config):
block = _ResidualBlock(
layers=num_layers,
num_features=_num_features,
in_channels=_in_channels,
out_channels=_out_channels,
dropout_prob=dropout_prob,
act=act,
norm=norm,
)
self.res_blocks.add_module(f"residualblock{i + 1}", block)
_in_channels = _out_channels
_out_channels *= 2
_num_features *= 2
# bottleneck convolution
self.bottleneck = nn.Sequential()
self.bottleneck.add_module(
"conv_bottleneck",
conv_type(_in_channels,
_num_features,
kernel_size=1,
stride=1,
padding=0,
bias=False))
self.upsample = UpSample(2,
scale_factor=2,
mode=UpsampleMode.NONTRAINABLE,
interp_mode=InterpolateMode.BILINEAR,
bias=False)
# decode branches
self.nucleus_prediction = _DecoderBranch(kernel_size=_ksize)
self.horizontal_vertical = _DecoderBranch(kernel_size=_ksize)
self.type_prediction: _DecoderBranch = None # type: ignore
if out_classes > 0:
self.type_prediction = _DecoderBranch(out_channels=out_classes,
kernel_size=_ksize)
for m in self.modules():
if isinstance(m, conv_type):
nn.init.kaiming_normal_(torch.as_tensor(m.weight))
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(torch.as_tensor(m.weight), 1)
nn.init.constant_(torch.as_tensor(m.bias), 0)
def __init__(
self,
num_features: int,
in_channels: int,
out_channels: int,
dropout_prob: float = 0.0,
act: Union[str, tuple] = ("relu", {
"inplace": True
}),
norm: Union[str, tuple] = "batch",
drop_first_norm_relu: int = 0,
kernel_size: int = 3,
) -> None:
"""Dense Convolutional Block.
References:
Huang, Gao, et al. "Densely connected convolutional networks."
Proceedings of the IEEE conference on computer vision and
pattern recognition. 2017.
Args:
num_features: number of internal channels used for the layer
in_channels: number of the input channels.
out_channels: number of the output channels.
dropout_prob: dropout rate after each dense layer.
act: activation type and arguments. Defaults to relu.
norm: feature normalization type and arguments. Defaults to batch norm.
drop_first_norm_relu - omits the first norm/relu for the first layer
kernel_size: size of the kernel for >1 convolutions (dependent on mode)
"""
super().__init__()
self.layers = nn.Sequential()
conv_type: Callable = Conv[Conv.CONV, 2]
dropout_type: Callable = Dropout[Dropout.DROPOUT, 2]
if not drop_first_norm_relu:
self.layers.add_module(
"preact_norm",
get_norm_layer(name=norm, spatial_dims=2,
channels=in_channels))
self.layers.add_module("preact_relu", get_act_layer(name=act))
self.layers.add_module(
"conv1",
conv_type(in_channels,
num_features,
kernel_size=1,
padding=0,
bias=False))
self.layers.add_module(
"norm2",
get_norm_layer(name=norm, spatial_dims=2, channels=num_features))
self.layers.add_module("relu2", get_act_layer(name=act))
if in_channels != 64 and drop_first_norm_relu:
self.layers.add_module(
"conv2",
conv_type(num_features,
num_features,
kernel_size=kernel_size,
stride=2,
padding=2,
bias=False))
else:
self.layers.add_module(
"conv2",
conv_type(num_features,
num_features,
kernel_size=1,
padding=0,
bias=False))
self.layers.add_module(
"norm3",
get_norm_layer(name=norm, spatial_dims=2, channels=num_features))
self.layers.add_module("relu3", get_act_layer(name=act))
self.layers.add_module(
"conv3",
conv_type(num_features,
out_channels,
kernel_size=1,
padding=0,
bias=False))
if dropout_prob > 0:
self.layers.add_module("dropout", dropout_type(dropout_prob))
def __init__(
self,
decode_config: Sequence[int] = (8, 4),
act: Union[str, tuple] = ("relu", {
"inplace": True
}),
norm: Union[str, tuple] = "batch",
dropout_prob: float = 0.0,
out_channels: int = 2,
kernel_size: int = 3,
) -> None:
"""
Args:
decode_config: number of layers for each block.
act: activation type and arguments. Defaults to relu.
norm: feature normalization type and arguments. Defaults to batch norm.
dropout_prob: dropout rate after each dense layer.
num_features: number of internal features used.
out_channels: number of the output channel.
kernel_size: size of the kernel for >1 convolutions (dependent on mode)
"""
super().__init__()
conv_type: Callable = Conv[Conv.CONV, 2]
# decode branches
_in_channels = 1024
_num_features = 128
_out_channels = 32
self.decoder_blocks = nn.Sequential()
for i, num_layers in enumerate(decode_config):
block = _DecoderBlock(
layers=num_layers,
num_features=_num_features,
in_channels=_in_channels,
out_channels=_out_channels,
dropout_prob=dropout_prob,
act=act,
norm=norm,
kernel_size=kernel_size,
)
self.decoder_blocks.add_module(f"decoderblock{i + 1}", block)
_in_channels = 512
# output layers
self.output_features = nn.Sequential()
_i = len(decode_config)
_pad_size = (kernel_size - 1) // 2
_seq_block = nn.Sequential(
OrderedDict([("conva",
conv_type(256,
64,
kernel_size=kernel_size,
stride=1,
bias=False,
padding=_pad_size))]))
self.output_features.add_module(f"decoderblock{_i + 1}", _seq_block)
_seq_block = nn.Sequential(
OrderedDict([
("norm", get_norm_layer(name=norm, spatial_dims=2,
channels=64)),
("relu", get_act_layer(name=act)),
("conv", conv_type(64, out_channels, kernel_size=1, stride=1)),
]))
self.output_features.add_module(f"decoderblock{_i + 2}", _seq_block)
self.upsample = UpSample(2,
scale_factor=2,
mode=UpsampleMode.NONTRAINABLE,
interp_mode=InterpolateMode.BILINEAR,
bias=False)
def __init__(
self,
in_channel: int,
out_channel: int,
kernel_size: int,
padding: int,
mode: int = 0,
act_name: Union[Tuple, str] = "RELU",
norm_name: Union[Tuple, str] = ("INSTANCE", {
"affine": True
}),
):
"""
Args:
in_channel: number of input channels.
out_channel: number of output channels.
kernel_size: kernel size to be expanded to 3D.
padding: padding size to be expanded to 3D.
mode: mode for the anisotropic kernels:
- 0: ``(k, k, 1)``, ``(1, 1, k)``,
- 1: ``(k, 1, k)``, ``(1, k, 1)``,
- 2: ``(1, k, k)``. ``(k, 1, 1)``.
act_name: activation layer type and arguments.
norm_name: feature normalization type and arguments.
"""
super().__init__()
self._in_channel = in_channel
self._out_channel = out_channel
self._p3dmode = int(mode)
conv_type = Conv[Conv.CONV, 3]
if self._p3dmode == 0: # (k, k, 1), (1, 1, k)
kernel_size0 = (kernel_size, kernel_size, 1)
kernel_size1 = (1, 1, kernel_size)
padding0 = (padding, padding, 0)
padding1 = (0, 0, padding)
elif self._p3dmode == 1: # (k, 1, k), (1, k, 1)
kernel_size0 = (kernel_size, 1, kernel_size)
kernel_size1 = (1, kernel_size, 1)
padding0 = (padding, 0, padding)
padding1 = (0, padding, 0)
elif self._p3dmode == 2: # (1, k, k), (k, 1, 1)
kernel_size0 = (1, kernel_size, kernel_size)
kernel_size1 = (kernel_size, 1, 1)
padding0 = (0, padding, padding)
padding1 = (padding, 0, 0)
else:
raise ValueError("`mode` must be 0, 1, or 2.")
self.add_module("acti", get_act_layer(name=act_name))
self.add_module(
"conv",
conv_type(
in_channels=self._in_channel,
out_channels=self._in_channel,
kernel_size=kernel_size0,
stride=1,
padding=padding0,
groups=1,
bias=False,
dilation=1,
),
)
self.add_module(
"conv_1",
conv_type(
in_channels=self._in_channel,
out_channels=self._out_channel,
kernel_size=kernel_size1,
stride=1,
padding=padding1,
groups=1,
bias=False,
dilation=1,
),
)
self.add_module(
"norm",
get_norm_layer(name=norm_name,
spatial_dims=3,
channels=self._out_channel))
