def __init__(
self,
in_channels: int,
out_channels: int,
kernel_sizes: Tuple[int, ...] = (5, 9, 13),
activation: Union[str, None] = "ReLU",
) -> None:
"""Initialize Spatical Pyramid Pooling module.
Args:
in_channels: number of incoming channels
out_channels: number of outgoing channels
kernel_sizes: kernel sizes to use
activation: Name of the activation to use in convolution.
"""
super().__init__()
hidden_channels = in_channels // 2
self.conv1 = Conv(in_channels,
hidden_channels,
1,
1,
activation=activation)
self.conv2 = Conv(
hidden_channels * (len(kernel_sizes) + 1),
out_channels,
1,
1,
activation=activation,
)
self.pooling_modules = nn.ModuleList([
nn.MaxPool2d(kernel_size=k_size, stride=1, padding=k_size // 2)
for k_size in kernel_sizes
])
def __init__(
self,
in_channels: int,
out_channels: int,
kernel_size: int = 5,
activation: Union[str, None] = "ReLU",
) -> None:
"""Initialize SPPF module.
Args:
in_channels: number of incoming channels
out_channels: number of outgoing channels
kernel_size: kernel sizes to use. 5 is equivalent to kernel_sizes=(5, 9, 13) in SPP
activation: Name of the activation to use in convolution.
"""
super().__init__()
hidden_channels = in_channels // 2
self.conv1 = Conv(in_channels,
hidden_channels,
1,
1,
activation=activation)
self.conv2 = Conv(hidden_channels * 4, out_channels, 1, 1)
self.pooling_module = nn.MaxPool2d(kernel_size=kernel_size,
stride=1,
padding=kernel_size // 2)
def __init__(
self,
in_channels: int,
out_channels: int,
shortcut=True,
groups: int = 1,
expansion: float = 0.5,
activation: Union[str, None] = "ReLU",
) -> None:
"""Initialize Bottleneck instance.
Args:
in_channels: number of incoming channels
out_channels: number of outgoing channels
shortcut: whether to use shortcut connection.
This only works when in_channels and out_channels are identical.
groups: number of group convolution number.
expansion: expansion ratio.
activation: Name of the activation to use in convolution.
"""
super().__init__()
expansion_channel = int(out_channels * expansion)
self.conv1 = Conv(in_channels, expansion_channel, 1, 1, activation=activation)
self.conv2 = Conv(
expansion_channel, out_channels, 3, 1, groups=groups, activation=activation
)
self.shortcut = shortcut and in_channels == out_channels
def __init__(
self,
in_channels: int,
out_channels: int,
stride: int = 1,
expand_ratio: int = 4,
activation: Union[str, None] = "ReLU",
) -> None:
"""MobileNet v2 block.
https://github.com/pytorch/vision/blob/972ca657416c5896ba6e0f5fe7c0d0f3e99e1087/torchvision/models/mobilenetv2.py
Args:
in_channels: number of incoming channels
out_channels: number of outgoing channels
stride: stride.
expand_ratio: expansion ratio.
activation: Name of the activation to use in convolution.
"""
super().__init__()
self.stride = stride
assert stride in [1, 2]
hidden_dim = int(round(in_channels * expand_ratio))
self.use_res_connect = self.stride == 1 and in_channels == out_channels
layers: List[nn.Module] = []
if expand_ratio != 1:
# pw
layers.append(
Conv(in_channels, hidden_dim, kernel_size=1, activation=activation)
)
layers.extend(
[
# dw
Conv(
hidden_dim,
hidden_dim,
kernel_size=3,
stride=stride,
groups=hidden_dim,
activation=activation,
),
# pw-linear
# TODO(jeikeilim): Add fuse forward here.
nn.Conv2d(hidden_dim, out_channels, 1, 1, 0, bias=False),
nn.BatchNorm2d(out_channels),
]
)
self.conv = nn.Sequential(*layers)
self.out_channels = out_channels
def __init__(
self,
in_channels: int,
out_channels: int,
n_repeat: int = 1,
shortcut=True,
groups: int = 1,
expansion: float = 0.5,
activation: Union[str, None] = "ReLU",
) -> None:
"""Initialize BottleneckCSP instance.
Args:
in_channels: number of incoming channels
out_channels: number of outgoing channels
n_repeat: repeat number of bottleneck.
shortcut: whether to use shortcut connection.
This only works when in_channels and out_channels are identical.
groups: number of group convolution number.
expansion: expansion ratio.
activation: Name of the activation to use in convolution.
"""
super().__init__()
expansion_channel = int(out_channels * expansion)
self.conv1 = Conv(in_channels, expansion_channel, 1, 1, activation=activation)
self.conv2 = nn.Conv2d(in_channels, expansion_channel, 1, 1, bias=False)
self.conv3 = nn.Conv2d(expansion_channel, expansion_channel, 1, 1, bias=False)
self.conv4 = Conv(
2 * expansion_channel, out_channels, 1, 1, activation=activation
)
self.batch_norm = nn.BatchNorm2d(2 * expansion_channel)
self.activation = Activation(activation)()
self.bottleneck_csp = nn.Sequential(
*[
Bottleneck(
expansion_channel,
expansion_channel,
shortcut=shortcut,
groups=groups,
expansion=1.0,
activation=activation,
)
for _ in range(n_repeat)
]
)
def __init__(
self,
in_channels: int,
conv_channels: int,
mlp_channels: int,
depth: int,
kernel_size: int = 3,
patch_size: Union[int, Tuple[int, int]] = 2,
dropout: float = 0.0,
activation: Union[str, None] = "SiLU",
) -> None:
"""Initialize Mobile ViT Block.
Args:
in_channels: number of incoming channels
conv_channels: number of channels to use in convolution.
mlp_channels: number of channels to use in MLP.
depth: depth of the transformer.
kernel_size: kernel size in nxn convolution.
dropout: dropout probability.
activation: Name of the activation to use in the middle of Linear modules.
"""
super().__init__()
self.patch_w, self.patch_h = ((patch_size, patch_size) if isinstance(
patch_size, int) else patch_size)
self.conv1_nxn = Conv(in_channels,
in_channels,
kernel_size,
activation=activation)
self.conv2_1x1 = Conv(in_channels,
conv_channels,
1,
activation=activation)
self.transformer = Transformer(conv_channels,
depth,
1,
32,
mlp_channels,
dropout=dropout)
self.conv3_1x1 = Conv(conv_channels,
in_channels,
1,
activation=activation)
self.conv4_nxn = Conv(2 * in_channels,
in_channels,
kernel_size,
activation=activation)
def __init__(
self,
in_channels: int,
out_channels: int,
shortcut=True,
groups: int = 1,
expansion: float = 0.5,
activation: Union[str, None] = "ReLU",
) -> None:
"""Initialize."""
super().__init__()
expansion_channel = int(out_channels * expansion)
self.conv1 = Conv(in_channels,
expansion_channel,
1,
1,
activation=activation)
self.conv2 = Conv(expansion_channel, out_channels, 3, 1, groups=groups)
self.shortcut = shortcut and in_channels == out_channels