def __init__(self,
num_classes,
init_conv_channels,
init_conv_stride,
channels,
mid_channels,
final_conv_channels,
strides,
in_channels=3,
in_size=(224, 224),
use_data_batchnorm=True):
super().__init__()
self.use_data_batchnorm = use_data_batchnorm
self.in_size = in_size
self.backbone = torch.nn.Sequential()
# data batchnorm
if self.use_data_batchnorm:
self.backbone.add_module(
"data_bn", torch.nn.BatchNorm2d(num_features=in_channels))
# init conv
self.backbone.add_module(
"init_conv",
conv3x3_block(in_channels=in_channels,
out_channels=init_conv_channels,
stride=init_conv_stride,
activation="relu6"))
# stages
in_channels = init_conv_channels
for stage_id, stage_channels in enumerate(channels):
stage = torch.nn.Sequential()
for unit_id, unit_channels in enumerate(stage_channels):
stride = strides[stage_id] if unit_id == 0 else 1
mid_channel = mid_channels[stage_id][unit_id]
stage.add_module(
"unit{}".format(unit_id + 1),
LinearBottleneck(in_channels=in_channels,
mid_channels=mid_channel,
out_channels=unit_channels,
stride=stride))
in_channels = unit_channels
self.backbone.add_module("stage{}".format(stage_id + 1), stage)
self.backbone.add_module(
"final_conv",
conv1x1_block(in_channels=in_channels,
out_channels=final_conv_channels,
activation="relu6"))
self.backbone.add_module("global_pool",
torch.nn.AdaptiveAvgPool2d(output_size=1))
# classifier
self.classifier = Classifier(in_channels=final_conv_channels,
num_classes=num_classes)
self.init_params()
def __init__(self,
channels,
num_classes,
preact=False,
init_unit_channels=64,
use_init_unit_large=True,
in_channels=3,
in_size=(224, 224),
use_data_batchnorm=True):
super().__init__()
self.use_data_batchnorm = use_data_batchnorm
self.use_init_unit_large = use_init_unit_large
self.in_size = in_size
self.backbone = torch.nn.Sequential()
# data batchnorm
if self.use_data_batchnorm:
self.backbone.add_module("data_bn", torch.nn.BatchNorm2d(num_features=in_channels))
# init unit
if self.use_init_unit_large:
self.backbone.add_module("init_unit", InitUnitLarge(in_channels=in_channels, out_channels=init_unit_channels))
else:
self.backbone.add_module("init_unit", InitUnitSmall(in_channels=in_channels, out_channels=init_unit_channels, preact=preact))
# stages
in_channels = init_unit_channels
for stage_id, stage_channels in enumerate(channels):
stage = torch.nn.Sequential()
for unit_id, unit_channels in enumerate(stage_channels):
stride = 2 if (unit_id == 0) and (stage_id != 0) else 1
if preact:
stage.add_module("unit{}".format(unit_id + 1), PreactUnit(in_channels=in_channels, out_channels=unit_channels, stride=stride))
else:
stage.add_module("unit{}".format(unit_id + 1), StandardUnit(in_channels=in_channels, out_channels=unit_channels, stride=stride))
in_channels = unit_channels
self.backbone.add_module("stage{}".format(stage_id + 1), stage)
if preact:
self.backbone.add_module("final_activation", PostActivation(in_channels))
self.backbone.add_module("global_pool", torch.nn.AdaptiveAvgPool2d(output_size=1))
# classifier
self.classifier = Classifier(in_channels=in_channels, num_classes=num_classes)
self.init_params()
class ResNet(torch.nn.Module):
def __init__(self,
channels,
num_classes,
preact=False,
init_unit_channels=64,
use_init_unit_large=True,
in_channels=3,
in_size=(224, 224),
use_data_batchnorm=True):
super().__init__()
self.use_data_batchnorm = use_data_batchnorm
self.use_init_unit_large = use_init_unit_large
self.in_size = in_size
self.backbone = torch.nn.Sequential()
# data batchnorm
if self.use_data_batchnorm:
self.backbone.add_module("data_bn", torch.nn.BatchNorm2d(num_features=in_channels))
# init unit
if self.use_init_unit_large:
self.backbone.add_module("init_unit", InitUnitLarge(in_channels=in_channels, out_channels=init_unit_channels))
else:
self.backbone.add_module("init_unit", InitUnitSmall(in_channels=in_channels, out_channels=init_unit_channels, preact=preact))
# stages
in_channels = init_unit_channels
for stage_id, stage_channels in enumerate(channels):
stage = torch.nn.Sequential()
for unit_id, unit_channels in enumerate(stage_channels):
stride = 2 if (unit_id == 0) and (stage_id != 0) else 1
if preact:
stage.add_module("unit{}".format(unit_id + 1), PreactUnit(in_channels=in_channels, out_channels=unit_channels, stride=stride))
else:
stage.add_module("unit{}".format(unit_id + 1), StandardUnit(in_channels=in_channels, out_channels=unit_channels, stride=stride))
in_channels = unit_channels
self.backbone.add_module("stage{}".format(stage_id + 1), stage)
if preact:
self.backbone.add_module("final_activation", PostActivation(in_channels))
self.backbone.add_module("global_pool", torch.nn.AdaptiveAvgPool2d(output_size=1))
# classifier
self.classifier = Classifier(in_channels=in_channels, num_classes=num_classes)
self.init_params()
def init_params(self):
# backbone
for name, module in self.backbone.named_modules():
if isinstance(module, torch.nn.Conv2d):
torch.nn.init.kaiming_uniform_(module.weight)
if module.bias is not None:
torch.nn.init.constant_(module.bias, 0)
# classifier
self.classifier.init_params()
def forward(self, x):
x = self.backbone(x)
x = self.classifier(x)
return x
class MobileNetV1(torch.nn.Module):
def __init__(self,
num_classes,
init_conv_channels,
init_conv_stride,
channels,
strides,
in_channels=3,
in_size=(224, 224),
use_data_batchnorm=True):
super().__init__()
self.use_data_batchnorm = use_data_batchnorm
self.in_size = in_size
self.backbone = torch.nn.Sequential()
# data batchnorm
if self.use_data_batchnorm:
self.backbone.add_module(
"data_bn", torch.nn.BatchNorm2d(num_features=in_channels))
# init conv
self.backbone.add_module(
"init_conv",
conv3x3_block(in_channels=in_channels,
out_channels=init_conv_channels,
stride=init_conv_stride))
# stages
in_channels = init_conv_channels
for stage_id, stage_channels in enumerate(channels):
stage = torch.nn.Sequential()
for unit_id, unit_channels in enumerate(stage_channels):
stride = strides[stage_id] if unit_id == 0 else 1
stage.add_module(
"unit{}".format(unit_id + 1),
DepthwiseSeparableConvBlock(in_channels=in_channels,
out_channels=unit_channels,
stride=stride))
in_channels = unit_channels
self.backbone.add_module("stage{}".format(stage_id + 1), stage)
self.backbone.add_module("global_pool",
torch.nn.AdaptiveAvgPool2d(output_size=1))
# classifier
self.classifier = Classifier(in_channels=in_channels,
num_classes=num_classes)
self.init_params()
def init_params(self):
# backbone
for name, module in self.backbone.named_modules():
if isinstance(module, torch.nn.Conv2d):
torch.nn.init.kaiming_uniform_(module.weight)
if module.bias is not None:
torch.nn.init.constant_(module.bias, 0)
# classifier
self.classifier.init_params()
def forward(self, x):
x = self.backbone(x)
x = self.classifier(x)
return x
def __init__(self,
num_classes,
init_conv_channels,
init_conv_stride,
final_conv_channels,
final_conv_se,
channels,
mid_channels,
strides,
se_units,
kernel_sizes,
activations,
dropout_rate=0.0,
in_channels=3,
in_size=(224, 224),
use_data_batchnorm=True):
super().__init__()
self.use_data_batchnorm = use_data_batchnorm
self.in_size = in_size
self.dropout_rate = dropout_rate
self.backbone = torch.nn.Sequential()
# data batchnorm
if self.use_data_batchnorm:
self.backbone.add_module(
"data_bn", torch.nn.BatchNorm2d(num_features=in_channels))
# init conv
self.backbone.add_module(
"init_conv",
conv3x3_block(in_channels=in_channels,
out_channels=init_conv_channels,
stride=init_conv_stride,
activation="hswish"))
# stages
in_channels = init_conv_channels
for stage_id, stage_channels in enumerate(channels):
stage = torch.nn.Sequential()
for unit_id, unit_channels in enumerate(stage_channels):
stride = strides[stage_id] if unit_id == 0 else 1
mid_channel = mid_channels[stage_id][unit_id]
use_se = se_units[stage_id][unit_id] == 1
kernel_size = kernel_sizes[stage_id]
activation = activations[stage_id]
stage.add_module(
"unit{}".format(unit_id + 1),
LinearBottleneck(in_channels=in_channels,
mid_channels=mid_channel,
out_channels=unit_channels,
stride=stride,
activation=activation,
use_se=use_se,
kernel_size=kernel_size))
in_channels = unit_channels
self.backbone.add_module("stage{}".format(stage_id + 1), stage)
self.backbone.add_module(
"final_conv1",
conv1x1_block(in_channels=in_channels,
out_channels=final_conv_channels[0],
activation="hswish"))
in_channels = final_conv_channels[0]
if final_conv_se:
self.backbone.add_module(
"final_se",
SEUnit(channels=in_channels,
squeeze_factor=4,
squeeze_activation="relu",
excite_activation="hsigmoid"))
self.backbone.add_module("final_pool",
torch.nn.AdaptiveAvgPool2d(output_size=1))
if len(final_conv_channels) > 1:
self.backbone.add_module(
"final_conv2",
conv1x1_block(in_channels=in_channels,
out_channels=final_conv_channels[1],
activation="hswish",
use_bn=False))
in_channels = final_conv_channels[1]
if self.dropout_rate != 0.0:
self.backbone.add_module("final_dropout",
torch.nn.Dropout(dropout_rate))
# classifier
self.classifier = Classifier(in_channels=in_channels,
num_classes=num_classes)
self.init_params()
class MobileNetV3(torch.nn.Module):
"""
Class for constructing MobileNetsV3.
If you are in doubt, please use the high-level function `get_mobilenet` to
obtain ready-to-use models.
"""
def __init__(self,
num_classes,
init_conv_channels,
init_conv_stride,
final_conv_channels,
final_conv_se,
channels,
mid_channels,
strides,
se_units,
kernel_sizes,
activations,
dropout_rate=0.0,
in_channels=3,
in_size=(224, 224),
use_data_batchnorm=True):
super().__init__()
self.use_data_batchnorm = use_data_batchnorm
self.in_size = in_size
self.dropout_rate = dropout_rate
self.backbone = torch.nn.Sequential()
# data batchnorm
if self.use_data_batchnorm:
self.backbone.add_module(
"data_bn", torch.nn.BatchNorm2d(num_features=in_channels))
# init conv
self.backbone.add_module(
"init_conv",
conv3x3_block(in_channels=in_channels,
out_channels=init_conv_channels,
stride=init_conv_stride,
activation="hswish"))
# stages
in_channels = init_conv_channels
for stage_id, stage_channels in enumerate(channels):
stage = torch.nn.Sequential()
for unit_id, unit_channels in enumerate(stage_channels):
stride = strides[stage_id] if unit_id == 0 else 1
mid_channel = mid_channels[stage_id][unit_id]
use_se = se_units[stage_id][unit_id] == 1
kernel_size = kernel_sizes[stage_id]
activation = activations[stage_id]
stage.add_module(
"unit{}".format(unit_id + 1),
LinearBottleneck(in_channels=in_channels,
mid_channels=mid_channel,
out_channels=unit_channels,
stride=stride,
activation=activation,
use_se=use_se,
kernel_size=kernel_size))
in_channels = unit_channels
self.backbone.add_module("stage{}".format(stage_id + 1), stage)
self.backbone.add_module(
"final_conv1",
conv1x1_block(in_channels=in_channels,
out_channels=final_conv_channels[0],
activation="hswish"))
in_channels = final_conv_channels[0]
if final_conv_se:
self.backbone.add_module(
"final_se",
SEUnit(channels=in_channels,
squeeze_factor=4,
squeeze_activation="relu",
excite_activation="hsigmoid"))
self.backbone.add_module("final_pool",
torch.nn.AdaptiveAvgPool2d(output_size=1))
if len(final_conv_channels) > 1:
self.backbone.add_module(
"final_conv2",
conv1x1_block(in_channels=in_channels,
out_channels=final_conv_channels[1],
activation="hswish",
use_bn=False))
in_channels = final_conv_channels[1]
if self.dropout_rate != 0.0:
self.backbone.add_module("final_dropout",
torch.nn.Dropout(dropout_rate))
# classifier
self.classifier = Classifier(in_channels=in_channels,
num_classes=num_classes)
self.init_params()
def init_params(self):
# backbone
for name, module in self.backbone.named_modules():
if isinstance(module, torch.nn.Conv2d):
torch.nn.init.kaiming_uniform_(module.weight)
if module.bias is not None:
torch.nn.init.constant_(module.bias, 0)
# classifier
self.classifier.init_params()
def forward(self, x):
x = self.backbone(x)
x = self.classifier(x)
return x
class MobileNetV2(torch.nn.Module):
"""
Class for constructing MobileNetsV2.
If you are in doubt, please use the high-level function `get_mobilenet` to
obtain ready-to-use models.
"""
def __init__(self,
num_classes,
init_conv_channels,
init_conv_stride,
channels,
mid_channels,
final_conv_channels,
strides,
in_channels=3,
in_size=(224, 224),
use_data_batchnorm=True):
super().__init__()
self.use_data_batchnorm = use_data_batchnorm
self.in_size = in_size
self.backbone = torch.nn.Sequential()
# data batchnorm
if self.use_data_batchnorm:
self.backbone.add_module(
"data_bn", torch.nn.BatchNorm2d(num_features=in_channels))
# init conv
self.backbone.add_module(
"init_conv",
conv3x3_block(in_channels=in_channels,
out_channels=init_conv_channels,
stride=init_conv_stride,
activation="relu6"))
# stages
in_channels = init_conv_channels
for stage_id, stage_channels in enumerate(channels):
stage = torch.nn.Sequential()
for unit_id, unit_channels in enumerate(stage_channels):
stride = strides[stage_id] if unit_id == 0 else 1
mid_channel = mid_channels[stage_id][unit_id]
stage.add_module(
"unit{}".format(unit_id + 1),
LinearBottleneck(in_channels=in_channels,
mid_channels=mid_channel,
out_channels=unit_channels,
stride=stride))
in_channels = unit_channels
self.backbone.add_module("stage{}".format(stage_id + 1), stage)
self.backbone.add_module(
"final_conv",
conv1x1_block(in_channels=in_channels,
out_channels=final_conv_channels,
activation="relu6"))
self.backbone.add_module("global_pool",
torch.nn.AdaptiveAvgPool2d(output_size=1))
# classifier
self.classifier = Classifier(in_channels=final_conv_channels,
num_classes=num_classes)
self.init_params()
def init_params(self):
# backbone
for name, module in self.backbone.named_modules():
if isinstance(module, torch.nn.Conv2d):
torch.nn.init.kaiming_uniform_(module.weight)
if module.bias is not None:
torch.nn.init.constant_(module.bias, 0)
# classifier
self.classifier.init_params()
def forward(self, x):
x = self.backbone(x)
x = self.classifier(x)
return x