def __init__(self, in_channels: int, out_channels: int,
stride: int = 1,
expansion: int = 4,
base_width: int = 64,
cardinality: int = 1,
radix: int = 1):
super(Bottleneck, self).__init__()
self.radix = radix
width = int((out_channels / expansion) * (base_width / 64) * cardinality)
self.conv1 = ConvBnReLU2d(in_channels, width, 1)
self.conv2 = ConvBnReLU2d(width, width * radix, 3, padding=1, groups=cardinality * radix)
self.sa = SplitAttention(width, width * radix, 4, cardinality=cardinality, radix=radix)
self.pool = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1) if stride == 2 else nn.Identity()
self.conv3 = ConvBn2d(width, out_channels, 1)
self.activation = nn.ReLU(inplace=True)
self.downsample = (
nn.Sequential(
nn.AvgPool2d(stride) if stride != 1 else nn.Identity(),
ConvBn2d(in_channels, out_channels, 1),
)
if stride != 1 or in_channels != out_channels
else nn.Identity()
)
def __init__(self, in_channels, out_channels, expansion_channels,
kernel_size=3,
stride=1,
use_se=False):
super(GhostBottleneck, self).__init__()
self.conv1 = GhostModule(in_channels, expansion_channels)
self.conv2 = (
ConvBn2d(expansion_channels, expansion_channels, kernel_size,
padding=kernel_size // 2, stride=stride, groups=expansion_channels)
if stride != 1
else nn.Identity()
)
self.se = SEBlock(expansion_channels, expansion_channels) if use_se else nn.Identity()
self.conv3 = GhostModule(expansion_channels, out_channels, use_relu=False)
self.downsample = (
nn.Sequential(
ConvBn2d(in_channels, in_channels, kernel_size,
padding=kernel_size//2, stride=stride,
groups=in_channels),
ConvBn2d(in_channels, out_channels, 1),
)
if in_channels != out_channels or stride != 1
else None
)
def __init__(self,
in_channels,
out_channels,
stride=1,
dilation=1,
use_residual=True):
super(BasicBlock, self).__init__()
self.use_residual = use_residual
self.conv1 = ConvBnReLU2d(in_channels,
out_channels,
3,
padding=dilation,
stride=stride,
dilation=dilation)
self.conv2 = ConvBn2d(out_channels,
out_channels,
3,
padding=dilation,
dilation=dilation)
if self.use_residual:
self.downsample = (
ConvBn2d(in_channels, out_channels, 1, stride=stride) if
in_channels != out_channels or stride != 1 else nn.Identity())
self.activation = nn.ReLU(inplace=True)
def __init__(self,
in_channels,
out_channels,
stride=1,
dilation=1,
expansion=4,
use_residual=True):
super(Bottleneck, self).__init__()
self.use_residual = use_residual
width = out_channels // expansion
self.conv1 = ConvBnReLU2d(in_channels, width, 1)
self.conv2 = ConvBnReLU2d(width,
width,
3,
padding=dilation,
stride=stride,
dilation=dilation)
self.conv3 = ConvBn2d(width, out_channels, 1)
if self.use_residual:
self.downsample = (
ConvBn2d(in_channels, out_channels, 1, stride=stride) if
in_channels != out_channels or stride != 1 else nn.Identity())
self.activation = nn.ReLU(inplace=True)
def __init__(self,
in_channels: int,
out_channels: int,
kernel_size: int,
stride: int,
expansion: int,
activation: type = nn.Hardswish,
use_se: bool = True):
super(InvertedResidual, self).__init__()
# TODO: check if this is indeed correct
# this is the same as the implementation from
# https://github.com/d-li14/mobilenetv3.pytorch/
self.activation_after_se = use_se and expansion != 1
width = round_channels(expansion * in_channels)
self.conv1 = (ConvBnAct2d(in_channels, width, 1, activation=activation)
if expansion != 1 else nn.Identity())
self.conv2 = ConvBn2d(width,
width,
kernel_size,
padding=kernel_size // 2,
stride=stride,
groups=width)
self.act2 = activation()
self.se = (SqueezeExcitation(width, width)
if use_se else nn.Identity())
self.conv3 = ConvBn2d(width, out_channels, 1)
def __init__(self,
in_channels,
out_channels,
stride=1,
group_width=1,
se_reduction=4):
super(BottleneckY, self).__init__()
self.conv1 = ConvBnReLU2d(in_channels, out_channels, 1)
self.conv2 = ConvBnReLU2d(out_channels,
out_channels,
3,
padding=1,
stride=stride,
groups=out_channels //
min(out_channels, group_width))
self.se = SqueezeExcitation(out_channels,
out_channels,
mid_channels=round(in_channels /
se_reduction))
self.conv3 = ConvBn2d(out_channels, out_channels, 1)
self.downsample = (ConvBn2d(
in_channels, out_channels, 1, stride=stride) if stride != 1
or in_channels != out_channels else nn.Identity())
self.activation = nn.ReLU(inplace=True)
def __init__(self,
in_channels,
out_channels,
projection_ratio=4,
dropout_p=0.1):
super().__init__()
width = in_channels // projection_ratio
self.conv1 = ConvBnPReLU2d(in_channels, width, 1)
self.conv2 = nn.Sequential(
OrderedDict([
('conv',
nn.ConvTranspose2d(width,
width,
3,
stride=2,
padding=1,
output_padding=1,
bias=False)),
('bn', nn.BatchNorm2d(width)),
('relu', nn.PReLU()),
]))
self.conv3 = ConvBn2d(width, out_channels, 1)
self.dropout = nn.Dropout2d(p=dropout_p)
self.upsample = nn.ModuleDict({
'unpool':
nn.MaxUnpool2d(kernel_size=2, stride=2),
'conv':
ConvBn2d(in_channels, out_channels, 1),
})
self.activation = nn.PReLU()
def __init__(self, in_channels, out_channels, dilation=1, dropout_p=0.0):
super(SSnbtBlock, self).__init__()
if in_channels != out_channels:
raise ValueError("input and output channels must match")
channels = in_channels // 2
self.left = nn.Sequential(
ConvBnReLU2d(channels, channels, (3, 1), padding=(1, 0)),
ConvBnReLU2d(channels, channels, (1, 3), padding=(0, 1)),
ConvBnReLU2d(channels,
channels, (3, 1),
padding=(dilation, 0),
dilation=(dilation, 1)),
ConvBn2d(channels,
channels, (1, 3),
padding=(0, dilation),
dilation=(1, dilation)),
)
self.right = nn.Sequential(
ConvBnReLU2d(channels, channels, (3, 1), padding=(1, 0)),
ConvBnReLU2d(channels, channels, (1, 3), padding=(0, 1)),
ConvBnReLU2d(channels,
channels, (3, 1),
padding=(dilation, 0),
dilation=(dilation, 1)),
ConvBn2d(channels,
channels, (1, 3),
padding=(0, dilation),
dilation=(1, dilation)),
)
self.activation = nn.ReLU(inplace=True)
self.dropout = nn.Dropout2d(p=dropout_p)
def __init__(self,
in_channels,
out_channels,
stride=1,
expansion=4,
base_width=64,
cardinality=1):
super(Bottleneck, self).__init__()
width = int(
(out_channels / expansion) * (base_width / 64) * cardinality)
self.conv1 = ConvBnReLU2d(in_channels, width, 1)
self.conv2 = ConvBnReLU2d(width,
width,
3,
padding=1,
stride=stride,
groups=cardinality)
self.conv3 = ConvBn2d(width, out_channels, 1)
self.downsample = (
ConvBn2d(in_channels, out_channels, 1, stride=stride)
if in_channels != out_channels or stride != 1 else nn.Identity())
self.activation = nn.ReLU(inplace=True)
def Classifier(in_channels, out_channels):
return nn.Sequential(
ConvBn2d(in_channels, in_channels, 3, padding=1, groups=in_channels),
ConvBnReLU2d(in_channels, in_channels, 1),
ConvBn2d(in_channels, in_channels, 3, padding=1, groups=in_channels),
ConvBnReLU2d(in_channels, in_channels, 1),
nn.Dropout(0.1),
nn.Conv2d(in_channels, out_channels, kernel_size=1),
)
def __init__(self, in_channels, out_channels):
super(FeatureFusionModule, self).__init__()
lowres_channels, highres_channels = in_channels
self.lowres = nn.Sequential(
ConvBnReLU2d(lowres_channels,
lowres_channels,
kernel_size=3,
padding=4,
dilation=4,
groups=lowres_channels),
ConvBn2d(lowres_channels, out_channels, 1))
self.highres = ConvBn2d(highres_channels, out_channels, 1)
def __init__(self, in_channels: int, out_channels: int, stride: int = 1):
super(BasicBlock, self).__init__()
self.conv1 = ConvBnReLU2d(in_channels,
out_channels,
3,
padding=1,
stride=stride)
self.conv2 = ConvBn2d(out_channels, out_channels, 3, padding=1)
self.downsample = (
ConvBn2d(in_channels, out_channels, 1, stride=stride)
if in_channels != out_channels or stride != 1 else nn.Identity())
self.activation = nn.ReLU(inplace=True)
def __init__(self, in_channels, out_channels, scale_factor):
super().__init__()
lowres_channels, highres_channels = in_channels
self.lowres = nn.Sequential(
ConvBnReLU2d(lowres_channels,
lowres_channels,
3,
padding=scale_factor,
dilation=scale_factor,
groups=lowres_channels),
ConvBn2d(lowres_channels, out_channels, 1),
)
self.highres = ConvBn2d(highres_channels, out_channels, 1)
def __init__(self, in_channels, out_channels):
super(Downsampling, self).__init__()
self.conv = ConvBn2d(in_channels, out_channels - in_channels,
kernel_size=3, padding=1, stride=2)
self.pool = nn.MaxPool2d(kernel_size=2, ceil_mode=True)
self.relu = nn.ReLU(inplace=True)
def __init__(self,
in_channels,
out_channels,
kernel_size,
stride=1,
expansion=6,
reduction=4,
dropout_p=0.2):
super(MobileInvertedBottleneck, self).__init__()
hidden_channels = in_channels * expansion
self.conv1 = (ConvBnSwish2d(in_channels, hidden_channels, 1)
if expansion != 1 else nn.Identity())
self.conv2 = ConvBnSwish2d(hidden_channels,
hidden_channels,
kernel_size,
padding=kernel_size // 2,
stride=stride,
groups=hidden_channels)
self.se = SqueezeExcitation(hidden_channels,
hidden_channels,
reduction=reduction * expansion)
self.conv3 = ConvBn2d(hidden_channels, out_channels, 1)
self.dropout = nn.Dropout2d(p=dropout_p)
def __init__(self,
in_channels: int,
out_channels: int,
kernel_sizes: Union[int, List[int]],
stride: int = 1,
expansion: int = 6,
groups: Tuple[int, int] = (1, 1),
use_se: bool = True,
se_reduction: int = 4,
activation: str = 'swish'):
super(InvertedResidual, self).__init__()
assert not (activation == 'relu' and use_se), \
"can only use the SE block if activation is swish"
width = in_channels * expansion
self.conv1 = ((ConvBnSwish2d if activation == 'swish' else
ConvBnReLU2d)(in_channels, width, 1, groups=groups[0])
if expansion != 1 else nn.Identity())
self.conv2 = (MixConvBnSwish2d if activation == 'swish' else
MixConvBnReLU2d)(width,
width,
kernel_sizes,
stride=stride)
self.se = SqueezeExcitation(
width, width, int(in_channels /
se_reduction)) if use_se else nn.Identity()
self.conv3 = ConvBn2d(width, out_channels, 1, groups=groups[1])
def __init__(self, in_channels, out_channels, stride=1, group_width=1):
super(BottleneckX, self).__init__()
self.downsample = (ConvBn2d(
in_channels, out_channels, 1, stride=stride) if stride != 1
or in_channels != out_channels else nn.Identity())
self.conv1 = ConvBnReLU2d(in_channels, out_channels, 1)
self.conv2 = ConvBnReLU2d(out_channels,
out_channels,
3,
padding=1,
stride=stride,
groups=max(out_channels // group_width, 1))
self.conv3 = ConvBn2d(out_channels, out_channels, 1)
self.activation = nn.ReLU(inplace=True)
def __init__(self,
in_channels: int,
out_channels: int,
stride: int = 1,
expansion: int = 4):
super(Bottleneck, self).__init__()
width = out_channels // expansion
self.conv1 = ConvBnReLU2d(in_channels, width, 1)
self.conv2 = ConvBnReLU2d(width, width, 3, padding=1, stride=stride)
self.conv3 = ConvBn2d(width, out_channels, 1)
self.downsample = (
ConvBn2d(in_channels, out_channels, 1, stride=stride)
if in_channels != out_channels or stride != 1 else nn.Identity())
self.activation = nn.ReLU(inplace=True)
def __init__(self, in_channels, out_channels, stride=1, expansion=6):
super(BottleneckBlock, self).__init__()
expansion_channels = in_channels * expansion
self.conv1 = ConvBnReLU2d(in_channels, expansion_channels, 1)
self.conv2 = ConvBnReLU2d(expansion_channels,
expansion_channels,
3,
padding=1,
stride=stride,
groups=expansion_channels)
self.conv3 = ConvBn2d(expansion_channels, out_channels, 1)
def __init__(self, in_channels, out_channels, stride=1, expansion=6):
super().__init__()
expansion_channels = expansion * in_channels
self.conv1 = ConvBnReLU2d(in_channels, expansion_channels, 1)
self.conv2 = ConvBnReLU2d(expansion_channels,
expansion_channels,
3,
padding=1,
stride=stride,
groups=expansion_channels)
self.conv3 = ConvBn2d(expansion_channels, out_channels, 1)
def __init__(self, in_channels, out_channels):
super().__init__()
# The Learning to Downsample module
# It encodes low level features in a efficient way
# It is composed by three convolutional layers, where the first one is
# a regular conv and the other two are depthwise separable conv
# layers.
# The first convolutional layer is a regular conv because there is no
# advantage in using a ds conv in such small number of channels.
# All layers are a spatial kernel of 3x3 and have a stride of 2 for a total downsample of 8 times.
# Also, there is no nonlinearity between the depthwise and pointwise conv.
self.downsample = nn.Sequential(
ConvBnReLU2d(in_channels, 32, 3, padding=1, stride=1),
ConvBn2d(32, 32, 3, padding=1, stride=2, groups=32),
ConvBnReLU2d(32, 48, 1),
ConvBn2d(48, 48, 3, padding=1, stride=2, groups=48),
ConvBnReLU2d(48, 64, 1),
)
# The Global Feature Extractor module is aimed at capturing the global
# context for the task of image segmentation.
# This module directly takes the 1/8 downsampled output of the
# Learning to Downsample module, performs a feature encoding using the
# MobileNet bottleneck residual block and then performs a pyramid pooling
# at the end to aggregate the different region-based context information.
self.features = nn.Sequential(
BottleneckModule(64, 64, expansion=6, repeats=3, stride=2),
BottleneckModule(64, 96, expansion=6, repeats=3, stride=2),
BottleneckModule(96, 128, expansion=6, repeats=3, stride=1),
PyramidPoolingModule(128, 128))
# The Feature Fusion adds the low-resolution features from the
# Global Feature Encoder and the high-resolution features from the
# Learning to Downsample Module.
self.fusion = FeatureFusionModule((128, 64), 128, scale_factor=4)
# The classifier discriminates the classes from the features produced
# by fusion module.
self.classifier = Classifier(128, out_channels)
def __init__(self, in_channels, out_channels):
super().__init__()
assert out_channels > in_channels, \
"output channels must be greater than the input channels"
self.conv = ConvBn2d(in_channels,
out_channels - in_channels,
kernel_size=3,
padding=1,
stride=2)
self.pool = nn.MaxPool2d(kernel_size=2)
self.activation = nn.PReLU()
def __init__(self, in_channels, out_channels, kernel_size, dropout_p=0.01):
super().__init__()
self.conv1 = nn.Sequential(
ConvBnReLU2d(in_channels, out_channels, (kernel_size, 1), padding=(kernel_size//2, 0)),
ConvBnReLU2d(out_channels, out_channels, (1, kernel_size), padding=(0, kernel_size//2)),
)
self.conv2 = nn.Sequential(
ConvBnReLU2d(out_channels, out_channels, (kernel_size, 1), padding=(kernel_size//2, 0)),
ConvBn2d(out_channels, out_channels, (1, kernel_size), padding=(0, kernel_size//2)),
)
self.activation = nn.ReLU(inplace=True)
self.dropout = nn.Dropout2d(p=dropout_p)
def __init__(self,
in_channels,
out_channels,
kernel_size=3,
dilation=1,
projection_ratio=4,
dropout_p=0.1):
super().__init__()
width = in_channels // projection_ratio
self.conv1 = ConvBnPReLU2d(in_channels, width, 1)
self.conv2 = (ConvBnPReLU2d(
width, width, 3, padding=dilation, dilation=dilation)
if kernel_size == 3 else nn.Sequential(
ConvBn2d(width, width, (1, 5), padding=(0, 2)),
ConvBn2d(width, width, (5, 1), padding=(2, 0)),
nn.PReLU(),
))
self.conv3 = ConvBn2d(width, out_channels, 1)
self.dropout = nn.Dropout2d(p=dropout_p)
self.activation = nn.PReLU()
def __init__(self,
in_channels,
out_channels,
projection_ratio=4,
dropout_p=0.1):
super().__init__()
width = in_channels // projection_ratio
self.conv1 = ConvBnPReLU2d(in_channels, width, 1)
self.conv2 = ConvBnPReLU2d(width, width, 3, 1, stride=2)
self.conv3 = ConvBn2d(width, out_channels, 1)
self.dropout = nn.Dropout2d(p=dropout_p)
self.downsample = nn.MaxPool2d(kernel_size=2, return_indices=True)
self.activation = nn.PReLU()
def __init__(self,
in_channels: int,
out_channels: int,
stride: int = 1,
expansion: int = 6):
super(InvertedResidual, self).__init__()
hidden_channels = in_channels * expansion
self.conv1 = (ConvBnReLU62d(in_channels, hidden_channels, 1)
if expansion != 1 else nn.Identity())
self.conv2 = ConvBnReLU62d(hidden_channels,
hidden_channels,
3,
padding=1,
stride=stride,
groups=hidden_channels)
self.conv3 = ConvBn2d(hidden_channels, out_channels, 1)
def __init__(self, in_channels, out_channels, dilation_rates=(2, 5, 9), dropout_p=0.01):
super().__init__()
self.conv1 = nn.Sequential(
ConvBnReLU2d(in_channels, out_channels, (3, 1), padding=(1, 0)),
ConvBnReLU2d(out_channels, out_channels, (1, 3), padding=(0, 1)),
)
self.conv2 = nn.ModuleList(
nn.Sequential(
ConvBnReLU2d(out_channels, out_channels, (3, 1), padding=(dilation, 0), dilation=(dilation, 1)),
ConvBn2d(out_channels, out_channels, (1, 3), padding=(0, dilation), dilation=(1, dilation)),
)
for dilation in dilation_rates
)
self.activation = nn.ReLU(inplace=True)
self.dropout = nn.Dropout2d(p=dropout_p)