def __init__(self, n_dim: int, channel: int, reduction: int = 16):
"""
Squeeze and Excitation Layer
https://arxiv.org/abs/1709.01507
Parameters
----------
n_dim : int
dimensionality of convolution
channel : int
number of input channel
reduction : int
channel reduction factor
"""
super().__init__()
self.pool = PoolingNd('AdaptiveAvg', n_dim, 1)
self.fc = torch.nn.Sequential(
ConvNd(n_dim,
channel,
channel // reduction,
kernel_size=1,
bias=False), torch.nn.ReLU(inplace=True),
ConvNd(n_dim,
channel // reduction,
channel,
kernel_size=1,
bias=False), torch.nn.Sigmoid())
def __init__(self,
num_classes: int,
in_channels: int,
n_dim: int = 2,
pool_type: str = "Max"):
super().__init__()
self.features = torch.nn.Sequential(
ConvNd(n_dim, in_channels, 64, kernel_size=11, stride=4,
padding=2),
torch.nn.ReLU(inplace=True),
PoolingNd(pool_type, n_dim, kernel_size=3, stride=2),
ConvNd(n_dim, 64, 192, kernel_size=5, padding=2),
torch.nn.ReLU(inplace=True),
PoolingNd(pool_type, n_dim, kernel_size=3, stride=2),
ConvNd(n_dim, 192, 384, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
ConvNd(n_dim, 384, 256, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
ConvNd(n_dim, 256, 256, kernel_size=3, padding=1),
torch.nn.ReLU(inplace=True),
PoolingNd(pool_type, n_dim, kernel_size=3, stride=2),
)
self.avgpool = PoolingNd("AdaptiveAvg", n_dim, 6)
self.classifier = torch.nn.Sequential(
torch.nn.Dropout(),
torch.nn.Linear(256 * pow(6, n_dim), 4096),
torch.nn.ReLU(inplace=True),
torch.nn.Dropout(),
torch.nn.Linear(4096, 4096),
torch.nn.ReLU(inplace=True),
torch.nn.Linear(4096, num_classes),
)
def __init__(self,
num_input_features: int,
growth_rate: int,
bn_size: int,
drop_rate: float,
n_dim: int = 2,
norm_type: str = "Batch"):
super().__init__()
self.add_module('norm1', NormNd(norm_type, n_dim, num_input_features)),
self.add_module('relu1', torch.nn.ReLU(inplace=True)),
self.add_module(
'conv1',
ConvNd(n_dim,
num_input_features,
bn_size * growth_rate,
kernel_size=1,
stride=1,
bias=False)),
self.add_module('norm2', NormNd(norm_type, n_dim,
bn_size * growth_rate)),
self.add_module('relu2', torch.nn.ReLU(inplace=True)),
self.add_module(
'conv2',
ConvNd(n_dim,
bn_size * growth_rate,
growth_rate,
kernel_size=3,
stride=1,
padding=1,
bias=False)),
self.drop_rate = drop_rate
def _conv1x1(n_dim: int, in_channels: int, out_channels: int,
groups: int = 1) -> ConvNd:
"""
Creates a module containing the appropriate convolution
Parameters
----------
n_dim : int
dimensionality of the input
in_channels : int
number of input channels
out_channels : int
number of output channels
groups : int
number of convolution groups
Returns
-------
ConvNd
appropriate convolution
"""
return ConvNd(
n_dim,
in_channels,
out_channels,
kernel_size=1,
groups=groups,
stride=1)
def _conv1x1(in_planes: int,
out_planes: int,
stride: Union[int, Sequence[int]] = 1,
n_dim: int = 2) -> ConvNd:
"""1x1 convolution"""
return ConvNd(n_dim,
in_planes,
out_planes,
kernel_size=1,
stride=stride,
bias=False)
def _conv3x3(in_planes: int,
out_planes: int,
stride: Union[int, Sequence[int]] = 1,
n_dim: int = 2) -> ConvNd:
"""3x3 convolution with padding"""
return ConvNd(n_dim,
in_planes,
out_planes,
kernel_size=3,
stride=stride,
padding=1,
bias=False)
def __init__(self,
block: torch.nn.Module,
layers: Sequence[int],
num_classes: int,
in_channels: int,
zero_init_residual: bool = False,
norm_layer: str = "Batch",
n_dim: int = 2,
start_filts: int = 64):
super().__init__()
self.start_filts = start_filts
self.inplanes = copy.deepcopy(start_filts)
self.conv1 = ConvNd(n_dim,
in_channels,
self.inplanes,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = NormNd(norm_layer, n_dim, self.inplanes)
self.relu = torch.nn.ReLU(inplace=True)
self.maxpool = PoolingNd("Max",
n_dim=n_dim,
kernel_size=3,
stride=2,
padding=1)
num_layers = 0
for idx, _layers in enumerate(layers):
stride = 1 if idx == 0 else 2
planes = min(self.start_filts * pow(2, idx), self.start_filts * 8)
_local_layer = self._make_layer(block,
planes,
_layers,
stride=stride,
norm_layer=norm_layer,
n_dim=n_dim)
setattr(self, "layer%d" % (idx + 1), _local_layer)
num_layers += 1
self.num_layers = num_layers
self.avgpool = PoolingNd("AdaptiveAvg", n_dim, 1)
self.fc = torch.nn.Linear(self.inplanes, num_classes)
self.reset_weights(zero_init_residual=zero_init_residual)
def __init__(self,
num_input_features: int,
num_output_features: int,
n_dim: int = 2,
norm_type: str = "Batch"):
super().__init__()
self.add_module('norm', NormNd(norm_type, n_dim, num_input_features))
self.add_module('relu', torch.nn.ReLU(inplace=True))
self.add_module(
'conv',
ConvNd(n_dim,
num_input_features,
num_output_features,
kernel_size=1,
stride=1,
bias=False))
self.add_module('pool', PoolingNd("AdaptiveAvg", n_dim, output_size=2))
def _upconv2x2(n_dim: int, in_channels: int, out_channels: int,
mode: str = 'transpose') -> torch.nn.Module:
"""
Creates a module to perform upsampling
Parameters
----------
n_dim : int
the dimensionality of the input
in_channels : int
number of input channels
out_channels : int
number of output channels
mode : str
mode for upsampling
Returns
-------
torch.nn.Module
the moule performing the upsampling (either on a convolutional basis
or on an interpolation basis)
"""
if mode == 'transpose':
return ConvNd(
n_dim,
in_channels,
out_channels,
kernel_size=2,
stride=2,
transposed=True
)
else:
# out_channels is always going to be the same
# as in_channels
if n_dim == 2:
upsample_mode = "bilinear"
elif n_dim == 3:
upsample_mode = "trilinear"
else:
raise ValueError
return torch.nn.Sequential(
torch.nn.Upsample(mode=upsample_mode, scale_factor=2),
_conv1x1(n_dim, in_channels, out_channels))
def __init__(self,
in_planes,
out_planes,
kernel_size=3,
stride=1,
groups=1,
n_dim=2,
norm_type="Batch"):
padding = (kernel_size - 1) // 2
super().__init__(
ConvNd(n_dim,
in_planes,
out_planes,
kernel_size,
stride,
padding,
groups=groups,
bias=False), NormNd(norm_type, n_dim, out_planes),
torch.nn.ReLU6(inplace=True))
def make_layers(cfg: Sequence[Union[int, str]], in_channels: int,
norm_type: str = None, n_dim: int = 2,
pool_type: str = "Max") -> torch.nn.Sequential:
layers = []
for v in cfg:
if v == 'P':
layers += [PoolingNd(pool_type, n_dim, kernel_size=2,
stride=2)]
else:
_layers = [ConvNd(n_dim, in_channels, v, kernel_size=3,
padding=1)]
if norm_type is not None:
_layers.append(NormNd(norm_type, n_dim, v))
_layers.append(torch.nn.ReLU(inplace=True))
layers += _layers
in_channels = v
return torch.nn.Sequential(*layers)
def _conv3x3(n_dim: int, in_channels: int, out_channels: int,
stride: Union[int, tuple, list] = 1,
padding: Union[int, tuple, list] = 1,
bias: bool = True, groups: int = 1) -> ConvNd:
"""
functional interface to create appropriate convolutional layer
Parameters
----------
n_dim : int
the dimensionality of the input
in_channels : int
number of input channels
out_channels : int
number of output channels
stride : int or tuple or list
convolutional strides
padding : int or tuple or list
convolutional padding
bias : bool
whether to include a bias
groups : int
number of convolution groups
Returns
-------
ConvNd
assembled convolution
"""
return ConvNd(
n_dim,
in_channels,
out_channels,
kernel_size=3,
stride=stride,
padding=padding,
bias=bias,
groups=groups)
def __init__(self,
inp,
oup,
stride,
expand_ratio,
n_dim=2,
norm_type="Batch"):
super().__init__()
self.stride = stride
assert stride in [1, 2]
hidden_dim = int(round(inp * expand_ratio))
self.use_res_connect = self.stride == 1 and inp == oup
layers = []
if expand_ratio != 1:
# pw
layers.append(
_ConvNormReLU(inp,
hidden_dim,
kernel_size=1,
n_dim=n_dim,
norm_type=norm_type))
layers.extend([
# dw
_ConvNormReLU(hidden_dim,
hidden_dim,
stride=stride,
groups=hidden_dim,
n_dim=n_dim,
norm_type=norm_type),
# pw-linear
ConvNd(n_dim, hidden_dim, oup, 1, 1, 0, bias=False),
NormNd(norm_type, n_dim, oup)
])
self.conv = torch.nn.Sequential(*layers)
def __init__(self,
in_channels: int,
channels: int,
stride: int,
cardinality: int,
width: int,
n_dim: int,
norm_layer: str,
reduction: int = 16):
"""
Squeeze and Excitation ResNeXt Block
Parameters
----------
in_channels : int
number of input channels
stride : int
stride of 3x3 convolution layer
cardinality : int
number of convolution groups
width : int
width of resnext block
n_dim : int
dimensionality of convolutions
norm_layer : str
type of normalization layer
reduction : int
reduction for se layer
"""
super().__init__()
out_channels = channels * self.expansion
if cardinality == 1:
rc = channels
else:
width_ratio = channels * (width / self.start_filts)
rc = cardinality * math.floor(width_ratio)
self.conv_reduce = ConvNd(n_dim,
in_channels,
rc,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.bn_reduce = NormNd(norm_layer, n_dim, rc)
self.relu = torch.nn.ReLU(inplace=True)
self.conv_conv = ConvNd(n_dim,
rc,
rc,
kernel_size=3,
stride=stride,
padding=1,
groups=cardinality,
bias=False)
self.bn = NormNd(norm_layer, n_dim, rc)
self.conv_expand = ConvNd(n_dim,
rc,
out_channels,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.bn_expand = NormNd(norm_layer, n_dim, out_channels)
self.shortcut = torch.nn.Sequential()
if in_channels != out_channels or stride != 1:
self.shortcut.add_module(
'shortcut_conv',
ConvNd(n_dim,
in_channels,
out_channels,
kernel_size=1,
stride=stride,
padding=0,
bias=False))
self.shortcut.add_module('shortcut_bn',
NormNd(norm_layer, n_dim, out_channels))
self.selayer = _SELayer(n_dim, out_channels, reduction=reduction)
def __init__(self,
num_classes: int,
in_channels: int,
growth_rate: int = 32,
block_config: Sequence[int] = (6, 12, 24, 16),
num_init_features: int = 64,
bn_size: int = 4,
drop_rate: float = 0,
n_dim: int = 2,
pool_type: str = "Max",
norm_type: str = "Batch"):
super().__init__()
# First convolution
self.features = torch.nn.Sequential(
OrderedDict([
('conv0',
ConvNd(n_dim,
in_channels,
num_init_features,
kernel_size=7,
stride=2,
padding=3,
bias=False)),
('norm0', NormNd(norm_type, n_dim, num_init_features)),
('relu0', torch.nn.ReLU(inplace=True)),
('pool0',
PoolingNd(pool_type,
n_dim,
kernel_size=3,
stride=2,
padding=1)),
]))
# Each denseblock
num_features = num_init_features
for i, num_layers in enumerate(block_config):
block = _DenseBlock(num_layers=num_layers,
num_input_features=num_features,
bn_size=bn_size,
growth_rate=growth_rate,
drop_rate=drop_rate,
n_dim=n_dim,
norm_type=norm_type)
self.features.add_module('denseblock%d' % (i + 1), block)
num_features = num_features + num_layers * growth_rate
if i != len(block_config) - 1:
trans = _Transition(num_input_features=num_features,
num_output_features=num_features // 2,
n_dim=n_dim,
norm_type=norm_type)
self.features.add_module('transition%d' % (i + 1), trans)
num_features = num_features // 2
# Final norm
self.features.add_module('norm5', NormNd(norm_type, n_dim,
num_features))
self.pool = PoolingNd("AdaptiveAvg", n_dim, 1)
# Linear layer
self.classifier = torch.nn.Linear(num_features, num_classes)
# Official init from torch repo.
for m in self.modules():
if isinstance(m, ConvNd):
torch.nn.init.kaiming_normal_(m.conv.weight)
elif isinstance(m, NormNd):
if hasattr(m.norm, "weight") and m.norm.weight is not None:
torch.nn.init.constant_(m.norm.weight, 1)
if hasattr(m.norm, "bias") and m.norm.bias is not None:
torch.nn.init.constant_(m.norm.bias, 0)
elif isinstance(m, torch.nn.Linear):
torch.nn.init.constant_(m.bias, 0)
def __init__(self,
n_dim: int,
norm_layer: str,
in_channels: int,
start_filts: int,
mode: str = '7x7'):
"""
Defines different sequences of start convolutions
Parameters
----------
n_dim : int
dimensionality of convolutions
norm_layer : str
type of normlization layer
in_channels : int
number of input channels
start_filts : int
number of channels after first convolution
mode : str
either '7x7' for default configuration (7x7 conv) or 3x3 for
three consecutive convolutions as proposed in
https://arxiv.org/abs/1812.01187
"""
super().__init__()
self._in_channels = in_channels
self._start_filts = start_filts
self._mode = mode
if mode == '7x7':
self.convs = torch.nn.Sequential(*[
ConvNd(n_dim,
in_channels,
self._start_filts,
kernel_size=7,
stride=2,
padding=3,
bias=False),
NormNd(norm_layer, n_dim, self._start_filts)
])
elif mode == '3x3':
self.convs = torch.nn.Sequential(*[
ConvNd(n_dim,
in_channels,
self._start_filts,
kernel_size=3,
stride=2,
padding=1,
bias=False),
NormNd(norm_layer, n_dim, self._start_filts),
ConvNd(n_dim,
self._start_filts,
self._start_filts,
kernel_size=3,
stride=1,
padding=1,
bias=False),
NormNd(norm_layer, n_dim, self._start_filts),
ConvNd(n_dim,
self._start_filts,
self._start_filts,
kernel_size=3,
stride=1,
padding=1,
bias=False),
NormNd(norm_layer, n_dim, self._start_filts)
])
else:
raise ValueError('{} is not a supported mode!'.format(mode))