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