def generate_profiling_primitives( self, base_channels, mult_ratio, strides, acts=None, use_ses=None, primitive_type="mobilenet_v2_block", spatial_size=224, stem_stride=2, stem_type="conv_3x3", sample=None, as_dict=True, ): channels = [make_divisible(c * mult_ratio, 8) for c in base_channels] primitives = [] acts = acts or [ None, ] * len(strides) use_ses = use_ses or [ None, ] * len(strides) sizes = [round(spatial_size / stem_stride + 1e-3)] for s in strides: sizes += [round(sizes[-1] / s + 1e-3)] stem_prim = Prim(stem_type, spatial_size, 3, channels[0], stem_stride, True) first_fixed_prim = Prim(primitive_type, spatial_size / stem_stride, channels[0], channels[1], strides[0], True, kernel_size=3, activation=acts[0], use_se=use_ses[0], expansion=1) primitives += [stem_prim, first_fixed_prim] for w, k, stage, i in self._traverse_search_space(sample): primitives.append( Prim( primitive_type, sizes[stage + i], channels[stage + i], channels[stage + 1], 1 if i else strides[stage], True, kernel_size=k, activation=acts[stage], use_se=use_ses[stage], expansion=w, )) if self.fixed_primitives is not None: primitives += self.fixed_primitives primitives = list(set(primitives)) if as_dict: primitives = [dict(p._asdict()) for p in primitives] return primitives
def __init__( self, expansion, C, C_out, stride, kernel_sizes=(3, 5, 7), do_kernel_transform=True, affine=True, activation="relu", schedule_cfg=None, ): FlexibleBlock.__init__(self, schedule_cfg) self.activation = activation C_inner = make_divisible(C * expansion, 8) self.kernel_sizes = sorted(kernel_sizes) self.kernel_size = self.kernel_sizes[-1] self.do_kernel_transform = do_kernel_transform self.affine = affine inv_bottleneck = None if expansion != 1: inv_bottleneck = nn.Sequential( FlexiblePointLinear(C, C_inner, 1, 1, 0), FlexibleBatchNorm2d(C_inner, affine=affine), get_op(activation)(), ) depth_wise = nn.Sequential( FlexibleDepthWiseConv( C_inner, self.kernel_sizes, stride, do_kernel_transform=do_kernel_transform, ), FlexibleBatchNorm2d(C_inner, affine=affine), get_op(activation)(), ) point_linear = nn.Sequential( FlexiblePointLinear(C_inner, C_out, 1, 1, 0), FlexibleBatchNorm2d(C_out, affine=affine), ) super(FlexibleMobileNetV2Block, self).__init__( expansion, C, C_out, stride, self.kernel_size, affine, activation, inv_bottleneck, depth_wise, point_linear, ) self.reset_mask()
def __init__( self, expansion, C, C_out, stride, kernel_size, affine, activation="relu", inv_bottleneck=None, depth_wise=None, point_linear=None, ): super(MobileNetV2Block, self).__init__() self.expansion = expansion self.C = C self.C_out = C_out self.C_inner = make_divisible(C * expansion, 8) self.stride = stride self.kernel_size = kernel_size self.act_fn = get_op(activation)() self.inv_bottleneck = None if expansion != 1: self.inv_bottleneck = inv_bottleneck or nn.Sequential( nn.Conv2d(C, self.C_inner, 1, 1, 0, bias=False), nn.BatchNorm2d(self.C_inner), self.act_fn) self.depth_wise = depth_wise or nn.Sequential( nn.Conv2d(self.C_inner, self.C_inner, self.kernel_size, stride, padding=self.kernel_size // 2, bias=False), nn.BatchNorm2d(self.C_inner), self.act_fn) self.point_linear = point_linear or nn.Sequential( nn.Conv2d(self.C_inner, C_out, 1, 1, 0, bias=False), nn.BatchNorm2d(C_out)) self.shortcut = nn.Sequential() self.has_conv_shortcut = False if stride == 1 and C != C_out: self.has_conv_shortcut = True self.shortcut = nn.Sequential( nn.Conv2d(C, C_out, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(C_out), )
def rollout_to_primitives(cls, rollout, primitive_type, spatial_size, strides, base_channels, mult_ratio=1., stem_type="conv_3x3", stem_stride=2, **kwargs): acts = kwargs.get("acts", [None] * len(rollout.depth)) use_ses = kwargs.get("use_ses", [None] * len(rollout.depth)) stem_stride = kwargs.get("stem_stride", 2) sizes = [ round(spatial_size / stem_stride / reduce(lambda p, q: p * q, strides[:i])) for i in range(1, len(strides) + 1) ] channels = [make_divisible(mult_ratio * c, 8) for c in base_channels] primitives = [ Prim(stem_type, spatial_size, 3, channels[0], stem_stride, True) ] for i, (depth, size, s, c_in, c_out, act, se) in enumerate( zip( rollout.depth, sizes, strides, channels[:-1], channels[1:], acts, use_ses, )): for j, width, kernel in zip(range(depth), rollout.width[i], rollout.kernel[i]): if j > 0: c_in = c_out s = 1 primitives.append( Prim( primitive_type, size, c_in, c_out, s, True, kernel_size=kernel, activation=act, use_se=se, expansion=width, )) return primitives
def set_mask(self, expansion, kernel_size): mask = None if expansion is not None and expansion != self.expansion: filters = self.point_linear[0].weight.data mask = _get_channel_mask(filters, make_divisible(self.C * expansion, 8)) if self.inv_bottleneck: self.inv_bottleneck[0].set_mask(None, mask) self.inv_bottleneck[1].set_mask(mask) self.depth_wise[0].set_mask(mask, kernel_size) self.depth_wise[1].set_mask(mask) self.point_linear[0].set_mask(mask, None)
def __init__(self, expansion, C, C_out, stride, affine, kernel_size=3, relu6=False): super(MobileNetBlock, self).__init__() C_inner = self.C_inner = make_divisible(expansion * C, 8) self.stride = stride self.relu6 = relu6 self.activation = F.relu6 if self.relu6 else F.relu self.conv1 = nn.Conv2d(C, C_inner, kernel_size=1, stride=1, padding=0, bias=False) self.bn1 = nn.BatchNorm2d(C_inner, affine=affine) padding = int((kernel_size - 1) / 2) self.conv2 = nn.Conv2d(C_inner, C_inner, kernel_size=kernel_size, stride=stride, padding=padding, groups=C_inner, bias=False) self.bn2 = nn.BatchNorm2d(C_inner, affine=affine) self.conv3 = nn.Conv2d(C_inner, C_out, kernel_size=1, stride=1, padding=0, bias=False) self.bn3 = nn.BatchNorm2d(C_out, affine=affine) self.shortcut = nn.Sequential() self.has_conv_shortcut = False if stride == 1 and C != C_out: self.has_conv_shortcut = True self.shortcut = nn.Sequential( nn.Conv2d(C, C_out, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(C_out, affine=affine), )
def __init__( self, device, blocks=[1, 4, 4, 4, 4, 4], strides=[1, 2, 2, 2, 1, 2], expansions=[1, 6, 6, 6, 6, 6], layer_channels=[16, 16, 24, 40, 80, 112, 160, 960, 1280], mult_ratio=1.0, kernel_sizes=[3, 5, 7], do_kernel_transform=True, use_ses=[False, False, True, False, True, True], acts=["relu", "relu", "relu", "h_swish", "h_swish", "h_swish"], num_classes=10, block_type="mbv3_block", pretrained_path=None, stem_stride=2, schedule_cfg=None, ): super(MobileNetV3Arch, self).__init__( device, blocks, strides, expansions, layer_channels, mult_ratio, kernel_sizes, do_kernel_transform, num_classes, block_type, pretrained_path, schedule_cfg, ) self.block_initializer = FlexibleBlock.get_class_(block_type) self.channels = [ make_divisible(c * mult_ratio, 8) for c in layer_channels ] self.stem_stride = stem_stride self.stem = nn.Sequential( nn.Conv2d(3, self.channels[0], kernel_size=3, stride=self.stem_stride, padding=1, bias=False), nn.BatchNorm2d(self.channels[0]), get_op("h_swish")(), ) expect( blocks[0] == expansions[0] == 1, "The first conv layer should have single block and no expansion.", ValueError, ) self.mult_ratio = mult_ratio self.use_ses = use_ses self.acts = acts self.cells = [ self.make_stage( self.channels[0], self.channels[1], self.blocks[0], self.strides[0], self.expansions[0], [3], self.use_ses[0], self.acts[0], ) ] for i, depth in enumerate(self.blocks[1:], 1): self.cells.append( self.make_stage( self.channels[i], self.channels[i + 1], depth, self.strides[i], self.expansions[i], self.kernel_sizes, self.use_ses[i], self.acts[i], )) self.cells = nn.ModuleList(self.cells) self.conv_head = nn.Sequential( nn.Conv2d(self.channels[-3], self.channels[-2], 1, 1, 0, bias=False), nn.BatchNorm2d(self.channels[-2]), get_op("h_swish")(), ) self.conv_final = nn.Sequential( nn.Conv2d(self.channels[-2], self.channels[-1], 1, 1, 0, bias=False), get_op("h_swish")(), ) self.classifier = nn.Linear(self.channels[-1], num_classes) if self.pretrained_path: state_dict = torch.load(self.pretrained_path, "cpu") if state_dict["classifier.weight"].shape[0] != self.num_classes: del state_dict["classifier.weight"] del state_dict["classifier.bias"] self.logger.info(self.load_state_dict(state_dict, strict=False)) self.to(self.device)
def rollout_to_primitives( cls, rollout, primitive_type, spatial_size, strides, base_channels, mult_ratio=1., stem_type="conv_3x3", # stem_stride=2, **kwargs): acts = kwargs.get("acts", [None] * len(rollout.depth)) use_ses = kwargs.get("use_ses", [None] * len(rollout.depth)) stem_stride = kwargs.get("stem_stride", 2) sizes = [round(spatial_size / stem_stride + 1e-3)] for s in strides: sizes += [round(sizes[-1] / s + 1e-3)] in_sizes = [spatial_size] + sizes # import ipdb; ipdb.set_trace() channels = [make_divisible(mult_ratio * c, 8) for c in base_channels] primitives = [ Prim(stem_type, spatial_size, 3, channels[0], stem_stride, True) ] for i, (depth, in_size, size, s, c_in, c_out, act, se) in enumerate( zip( rollout.depth, in_sizes[1:], sizes[1:], strides, channels[:-1], channels[1:], acts, use_ses, )): # `size` indicates the output size of this block # if the stride equals 2, the input size of this block should be `size * s` for j, width, kernel in zip(range(depth), rollout.width[i], rollout.kernel[i]): # if j > 0: # c_in = c_out # s = 1 # print(">>>", size, s, size * s) primitives.append( Prim( primitive_type, in_size if j == 0 else size, c_in if j == 0 else c_out, c_out, s if j == 0 else 1, True, kernel_size=kernel, activation=act, use_se=se, expansion=width, )) return primitives