def export(self) -> MCUNet:
export_model = MCUNet.__new__(MCUNet)
nn.Module.__init__(export_model)
# input stem
input_stem = OpSequential(
[
self.backbone["input_stem"].op_list[0].export(),
ResidualBlock(
self.backbone["input_stem"].op_list[1].conv.export(),
self.backbone["input_stem"].op_list[1].shortcut,
),
]
)
# stages
stages = []
for stage in self.backbone["stages"]:
blocks = []
for block in stage.op_list:
blocks.append(
ResidualBlock(
block.conv.export(),
block.shortcut,
)
)
stages.append(OpSequential(blocks))
# head
head = OpSequential(
[
ResidualBlock(
self.head.op_list[0].conv.export(),
self.head.op_list[0].shortcut,
),
self.head.op_list[1],
self.head.op_list[2].export(),
]
)
export_model.backbone = nn.ModuleDict(
{
"input_stem": input_stem,
"stages": nn.ModuleList(stages),
}
)
export_model.head = head
return export_model
def __init__(self, channel_divisor=8, n_classes=1000, dropout_rate=0):
super(TinyMobileNetV2, self).__init__(0.35, channel_divisor, n_classes,
dropout_rate)
self.head = OpSequential([
ResidualBlock(
InvertedBlock(
56,
112,
3,
expand_ratio=6,
act_func=("relu6", "relu6", None),
),
shortcut=None,
),
ConvLayer(112, 448, 1, act_func="relu6"),
nn.AdaptiveAvgPool2d(1),
LinearLayer(448, n_classes, dropout_rate=dropout_rate),
])
def __init__(self,
width_mult=1.0,
channel_divisor=8,
n_classes=1000,
dropout_rate=0):
super(MobileNetV3, self).__init__()
stage_width_list = [16, 24, 40, 80, 112, 160]
head_width_list = [960, 1280]
block_configs = [
[[64, 72], 3, 2, 2, "relu", False],
[[72, 120, 120], 5, 3, 2, "relu", True],
[[240, 200, 184, 184], 3, 4, 2, "h_swish", False],
[[480, 672], 3, 2, 1, "h_swish", True],
[[672, 960, 960], 5, 3, 2, "h_swish", True],
]
for i, w in enumerate(stage_width_list):
stage_width_list[i] = make_divisible(w * width_mult,
channel_divisor)
for i, w in enumerate(head_width_list):
head_width_list[i] = make_divisible(w * width_mult,
channel_divisor)
head_width_list[1] = max(head_width_list[1], 1280)
input_stem = OpSequential([
ConvLayer(3, stage_width_list[0], 3, 2, act_func="h_swish"),
ResidualBlock(
DsConvLayer(
stage_width_list[0],
stage_width_list[0],
3,
1,
("relu", None),
),
shortcut=nn.Identity(),
),
])
# stages
stages = []
in_channels = stage_width_list[0]
for (mid_c_list, ks, n, s, act_func,
use_se), c in zip(block_configs, stage_width_list[1:]):
blocks = []
for i in range(n):
stride = s if i == 0 else 1
mid_channels = make_divisible(
round(mid_c_list[i] * width_mult), channel_divisor)
if use_se:
conv = SeInvertedBlock(
in_channels,
c,
ks,
stride,
mid_channels=mid_channels,
act_func=(act_func, act_func, None),
se_config={
"act_func":
"relu",
"mid_channels":
max(
make_divisible(mid_channels / 4,
channel_divisor), 16),
},
)
else:
conv = InvertedBlock(
in_channels,
c,
ks,
stride,
mid_channels=mid_channels,
act_func=(act_func, act_func, None),
)
mb_conv = ResidualBlock(
conv,
shortcut=nn.Identity() if
(stride == 1 and in_channels == c and i != 0) else None,
)
blocks.append(mb_conv)
in_channels = c
stages.append(OpSequential(blocks))
# head
head = OpSequential([
ConvLayer(in_channels, head_width_list[0], 1, act_func="h_swish"),
nn.AdaptiveAvgPool2d(1),
ConvLayer(
head_width_list[0],
head_width_list[1],
1,
act_func="h_swish",
norm=None,
use_bias=True,
),
LinearLayer(head_width_list[1],
n_classes,
dropout_rate=dropout_rate),
])
self.backbone = nn.ModuleDict({
"input_stem": input_stem,
"stages": nn.ModuleList(stages),
})
self.head = head
def __init__(
self,
base_net: MobileNetV3,
aug_expand_list: List[float],
aug_width_mult_list: List[float],
n_classes: int,
dropout_rate=0.0,
):
nn.Module.__init__(self)
max_width_mult = max(aug_width_mult_list)
# input stem
base_input_stem = base_net.backbone["input_stem"]
aug_input_stem = OpSequential([
DynamicConvLayer(
3,
aug_width(base_input_stem.op_list[0].out_channels,
aug_width_mult_list, 1),
stride=2,
act_func="h_swish",
),
ResidualBlock(
DynamicDsConvLayer(
make_divisible(
base_input_stem.op_list[1].conv.in_channels *
max_width_mult,
1,
),
aug_width(
base_input_stem.op_list[1].conv.out_channels,
aug_width_mult_list,
1,
),
act_func=("relu", None),
),
shortcut=nn.Identity(),
),
])
# stages
aug_stages = []
for base_stage in base_net.backbone["stages"]:
stage = []
for base_block in base_stage.op_list:
if isinstance(base_block.conv, SeInvertedBlock):
se_config = {
"reduction":
(base_block.conv.se_layer.in_channels /
base_block.conv.se_layer.mid_channels + 1.0e-10),
"act_func":
base_block.conv.se_layer.act,
}
dynamic_block_cls = partial(DynamicSeInvertedBlock,
se_config=se_config)
elif isinstance(base_block.conv, InvertedBlock):
dynamic_block_cls = DynamicInvertedBlock
else:
raise NotImplementedError
stage.append(
ResidualBlock(
dynamic_block_cls(
in_channels=make_divisible(
base_block.conv.in_channels * max_width_mult,
1),
out_channels=aug_width(
base_block.conv.out_channels,
aug_width_mult_list, 1),
kernel_size=base_block.conv.kernel_size,
expand_ratio=aug_width(
base_block.conv.expand_ratio, aug_expand_list),
stride=base_block.conv.stride,
act_func=(
base_block.conv.inverted_conv.act,
base_block.conv.depth_conv.act,
base_block.conv.point_conv.act,
),
),
shortcut=base_block.shortcut,
))
aug_stages.append(OpSequential(stage))
# head
base_head = base_net.head
aug_head = OpSequential([
DynamicConvLayer(
make_divisible(
base_head.op_list[0].in_channels * max_width_mult, 1),
aug_width(base_head.op_list[0].out_channels,
aug_width_mult_list, 1),
1,
act_func=base_head.op_list[0].act,
),
nn.AdaptiveAvgPool2d(1),
DynamicConvLayer(
make_divisible(
base_head.op_list[2].in_channels * max_width_mult, 1),
aug_width(base_head.op_list[2].out_channels,
aug_width_mult_list, 1),
1,
use_bias=True,
norm=None,
act_func=base_head.op_list[2].act,
),
DynamicLinearLayer(
make_divisible(
base_head.op_list[-1].in_features * max_width_mult, 1),
n_classes,
dropout_rate=dropout_rate,
),
])
self.backbone = nn.ModuleDict({
"input_stem": aug_input_stem,
"stages": nn.ModuleList(aug_stages),
})
self.head = aug_head
def __init__(self, channel_divisor=8, n_classes=1000, dropout_rate=0):
super(MCUNet, self).__init__()
stage_width_list = [16, 8, 16, 24, 40, 48, 96]
head_width_list = [160]
act_func = "relu6"
block_configs = [
[[3, 5, 5, 4], [7, 3, 7, 5], 4, 2],
[[5, 5, 5], [5, 5, 5], 3, 2],
[[5, 6, 4], [3, 7, 5], 3, 2],
[[5, 5, 5], [5, 7, 3], 3, 1],
[[6, 5, 4], [3, 7, 3], 3, 2],
]
input_stem = OpSequential([
ConvLayer(3, stage_width_list[0], 3, 2, act_func=act_func),
ResidualBlock(
DsConvLayer(
stage_width_list[0],
stage_width_list[1],
3,
1,
(act_func, None),
),
shortcut=None,
),
])
# stages
stages = []
in_channels = stage_width_list[1]
for (e_list, ks_list, n, s), c in zip(block_configs,
stage_width_list[2:]):
blocks = []
for i in range(n):
stride = s if i == 0 else 1
mid_channels = make_divisible(round(e_list[i] * in_channels),
channel_divisor)
mb_conv = ResidualBlock(
InvertedBlock(
in_channels,
c,
ks_list[i],
stride,
mid_channels=mid_channels,
act_func=(act_func, act_func, None),
),
shortcut=nn.Identity() if
(stride == 1 and in_channels == c and i != 0) else None,
)
blocks.append(mb_conv)
in_channels = c
stages.append(OpSequential(blocks))
# head
head = OpSequential([
ResidualBlock(
InvertedBlock(
in_channels,
head_width_list[0],
7,
mid_channels=480,
act_func=(act_func, act_func, None),
),
shortcut=None,
),
nn.AdaptiveAvgPool2d(1),
LinearLayer(head_width_list[0],
n_classes,
dropout_rate=dropout_rate),
])
self.backbone = nn.ModuleDict({
"input_stem": input_stem,
"stages": nn.ModuleList(stages),
})
self.head = head
def __init__(
self,
base_net: MobileNetV2,
aug_expand_list: List[float],
aug_width_mult_list: List[float],
n_classes: int,
dropout_rate=0.0,
):
nn.Module.__init__(self)
max_width_mult = max(aug_width_mult_list)
# input stem
base_input_stem = base_net.backbone["input_stem"]
aug_input_stem = OpSequential([
DynamicConvLayer(
3,
aug_width(base_input_stem.op_list[0].out_channels,
aug_width_mult_list, 1),
stride=2,
act_func="relu6",
),
ResidualBlock(
DynamicDsConvLayer(
make_divisible(
base_input_stem.op_list[0].out_channels *
max_width_mult, 1),
aug_width(
base_input_stem.op_list[1].conv.out_channels,
aug_width_mult_list,
1,
),
act_func=("relu6", None),
),
shortcut=None,
),
])
# stages
aug_stages = []
for base_stage in base_net.backbone["stages"]:
stage = []
for base_block in base_stage.op_list:
stage.append(
ResidualBlock(
DynamicInvertedBlock(
in_channels=make_divisible(
base_block.conv.in_channels * max_width_mult,
1),
out_channels=aug_width(
base_block.conv.out_channels,
aug_width_mult_list, 1),
kernel_size=base_block.conv.kernel_size,
expand_ratio=aug_width(
base_block.conv.expand_ratio, aug_expand_list),
stride=base_block.conv.stride,
act_func=(
base_block.conv.inverted_conv.act,
base_block.conv.depth_conv.act,
base_block.conv.point_conv.act,
),
),
shortcut=base_block.shortcut,
))
aug_stages.append(OpSequential(stage))
# head
base_head = base_net.head
aug_head = OpSequential([
ResidualBlock(
DynamicInvertedBlock(
make_divisible(
base_head.op_list[0].conv.in_channels * max_width_mult,
1),
aug_width(
base_head.op_list[0].conv.out_channels,
aug_width_mult_list,
1,
),
base_head.op_list[0].conv.kernel_size,
expand_ratio=aug_width(
base_head.op_list[0].conv.expand_ratio,
aug_expand_list),
act_func=("relu6", "relu6", None),
),
shortcut=None,
),
DynamicConvLayer(
make_divisible(
base_head.op_list[1].in_channels * max_width_mult, 1),
aug_width(base_head.op_list[1].out_channels,
aug_width_mult_list, 1),
1,
act_func=base_head.op_list[1].act,
),
nn.AdaptiveAvgPool2d(1),
DynamicLinearLayer(
make_divisible(
base_head.op_list[-1].in_features * max_width_mult, 1),
n_classes,
dropout_rate=dropout_rate,
),
])
self.backbone = nn.ModuleDict({
"input_stem": aug_input_stem,
"stages": nn.ModuleList(aug_stages),
})
self.head = aug_head
def __init__(self,
width_mult=1.0,
channel_divisor=8,
n_classes=1000,
dropout_rate=0):
super(MobileNetV2, self).__init__()
stage_width_list = [32, 16, 24, 32, 64, 96, 160]
head_width_list = [320, 1280]
act_func = "relu6"
block_configs = [
# t, n, s
[6, 2, 2],
[6, 3, 2],
[6, 4, 2],
[6, 3, 1],
[6, 3, 2],
]
for i, w in enumerate(stage_width_list):
stage_width_list[i] = make_divisible(w * width_mult,
channel_divisor)
for i, w in enumerate(head_width_list):
head_width_list[i] = make_divisible(w * width_mult,
channel_divisor)
head_width_list[1] = max(head_width_list[1], 1280)
input_stem = OpSequential([
ConvLayer(3, stage_width_list[0], 3, 2, act_func=act_func),
ResidualBlock(
DsConvLayer(
stage_width_list[0],
stage_width_list[1],
3,
1,
(act_func, None),
),
shortcut=None,
),
])
# stages
stages = []
in_channels = stage_width_list[1]
for (t, n, s), c in zip(block_configs, stage_width_list[2:]):
blocks = []
for i in range(n):
stride = s if i == 0 else 1
mid_channels = make_divisible(round(t * in_channels),
channel_divisor)
mb_conv = ResidualBlock(
InvertedBlock(
in_channels,
c,
3,
stride,
mid_channels=mid_channels,
act_func=(act_func, act_func, None),
),
shortcut=nn.Identity() if
(stride == 1 and in_channels == c and i != 0) else None,
)
blocks.append(mb_conv)
in_channels = c
stages.append(OpSequential(blocks))
# head
head = OpSequential([
ResidualBlock(
InvertedBlock(
in_channels,
head_width_list[0],
3,
expand_ratio=6,
act_func=(act_func, act_func, None),
),
shortcut=None,
),
ConvLayer(head_width_list[0],
head_width_list[1],
1,
act_func=act_func),
nn.AdaptiveAvgPool2d(1),
LinearLayer(head_width_list[1],
n_classes,
dropout_rate=dropout_rate),
])
self.backbone = nn.ModuleDict({
"input_stem": input_stem,
"stages": nn.ModuleList(stages),
})
self.head = head