| |
| |
| |
|
|
| import json |
| import torch.nn as nn |
|
|
| from proard.utils.layers import ( |
| set_layer_from_config, |
| MBConvLayer, |
| ConvLayer, |
| IdentityLayer, |
| LinearLayer, |
| ResidualBlock, |
| ) |
| from proard.utils import ( |
| download_url, |
| make_divisible, |
| val2list, |
| MyNetwork, |
| MyGlobalAvgPool2d, |
| ) |
|
|
| __all__ = ["proxyless_base_cifar","proxyless_base", "ProxylessNASNets", "MobileNetV2", "ProxylessNASNets_Cifar", "MobileNetV2_Cifar"] |
|
|
|
|
| def proxyless_base( |
| net_config=None, |
| n_classes=None, |
| bn_param=None, |
| dropout_rate=None, |
| local_path="~/.torch/proxylessnas/", |
| ): |
| assert net_config is not None, "Please input a network config" |
| if "http" in net_config: |
| net_config_path = download_url(net_config, local_path) |
| else: |
| net_config_path = net_config |
| net_config_json = json.load(open(net_config_path, "r")) |
|
|
| if n_classes is not None: |
| net_config_json["classifier"]["out_features"] = n_classes |
| if dropout_rate is not None: |
| net_config_json["classifier"]["dropout_rate"] = dropout_rate |
|
|
| net = ProxylessNASNets.build_from_config(net_config_json) |
| if bn_param is not None: |
| net.set_bn_param(*bn_param) |
|
|
| return net |
|
|
|
|
| class ProxylessNASNets(MyNetwork): |
| def __init__(self, first_conv, blocks, feature_mix_layer, classifier): |
| super(ProxylessNASNets, self).__init__() |
|
|
| self.first_conv = first_conv |
| self.blocks = nn.ModuleList(blocks) |
| self.feature_mix_layer = feature_mix_layer |
| self.global_avg_pool = MyGlobalAvgPool2d(keep_dim=False) |
| self.classifier = classifier |
|
|
| def forward(self, x): |
| x = self.first_conv(x) |
| for block in self.blocks: |
| x = block(x) |
| if self.feature_mix_layer is not None: |
| x = self.feature_mix_layer(x) |
| x = self.global_avg_pool(x) |
| x = self.classifier(x) |
| return x |
|
|
| @property |
| def module_str(self): |
| _str = self.first_conv.module_str + "\n" |
| for block in self.blocks: |
| _str += block.module_str + "\n" |
| _str += self.feature_mix_layer.module_str + "\n" |
| _str += self.global_avg_pool.__repr__() + "\n" |
| _str += self.classifier.module_str |
| return _str |
|
|
| @property |
| def config(self): |
| return { |
| "name": ProxylessNASNets.__name__, |
| "bn": self.get_bn_param(), |
| "first_conv": self.first_conv.config, |
| "blocks": [block.config for block in self.blocks], |
| "feature_mix_layer": None |
| if self.feature_mix_layer is None |
| else self.feature_mix_layer.config, |
| "classifier": self.classifier.config, |
| } |
|
|
| @staticmethod |
| def build_from_config(config): |
| first_conv = set_layer_from_config(config["first_conv"]) |
| feature_mix_layer = set_layer_from_config(config["feature_mix_layer"]) |
| classifier = set_layer_from_config(config["classifier"]) |
|
|
| blocks = [] |
| for block_config in config["blocks"]: |
| blocks.append(ResidualBlock.build_from_config(block_config)) |
|
|
| net = ProxylessNASNets(first_conv, blocks, feature_mix_layer, classifier) |
| if "bn" in config: |
| net.set_bn_param(**config["bn"]) |
| else: |
| net.set_bn_param(momentum=0.1, eps=1e-3) |
|
|
| return net |
|
|
| def zero_last_gamma(self): |
| for m in self.modules(): |
| if isinstance(m, ResidualBlock): |
| if isinstance(m.conv, MBConvLayer) and isinstance( |
| m.shortcut, IdentityLayer |
| ): |
| m.conv.point_linear.bn.weight.data.zero_() |
|
|
| @property |
| def grouped_block_index(self): |
| info_list = [] |
| block_index_list = [] |
| for i, block in enumerate(self.blocks[1:], 1): |
| if block.shortcut is None and len(block_index_list) > 0: |
| info_list.append(block_index_list) |
| block_index_list = [] |
| block_index_list.append(i) |
| if len(block_index_list) > 0: |
| info_list.append(block_index_list) |
| return info_list |
|
|
| def load_state_dict(self, state_dict, **kwargs): |
| current_state_dict = self.state_dict() |
|
|
| for key in state_dict: |
| if key not in current_state_dict: |
| assert ".mobile_inverted_conv." in key |
| new_key = key.replace(".mobile_inverted_conv.", ".conv.") |
| else: |
| new_key = key |
| current_state_dict[new_key] = state_dict[key] |
| super(ProxylessNASNets, self).load_state_dict(current_state_dict) |
|
|
|
|
| class MobileNetV2(ProxylessNASNets): |
| def __init__( |
| self, |
| n_classes=1000, |
| width_mult=1.0, |
| bn_param=(0.1, 1e-3), |
| dropout_rate=0.2, |
| ks=None, |
| expand_ratio=None, |
| depth_param=None, |
| stage_width_list=None, |
| ): |
|
|
| ks = 3 if ks is None else ks |
| expand_ratio = 6 if expand_ratio is None else expand_ratio |
|
|
| input_channel = 32 |
| last_channel = 1280 |
|
|
| input_channel = make_divisible( |
| input_channel * width_mult, MyNetwork.CHANNEL_DIVISIBLE |
| ) |
| last_channel = ( |
| make_divisible(last_channel * width_mult, MyNetwork.CHANNEL_DIVISIBLE) |
| if width_mult > 1.0 |
| else last_channel |
| ) |
|
|
| inverted_residual_setting = [ |
| |
| [1, 16, 1, 1], |
| [expand_ratio, 24, 2, 2], |
| [expand_ratio, 32, 3, 2], |
| [expand_ratio, 64, 4, 2], |
| [expand_ratio, 96, 3, 1], |
| [expand_ratio, 160, 3, 2], |
| [expand_ratio, 320, 1, 1], |
| ] |
|
|
| if depth_param is not None: |
| assert isinstance(depth_param, int) |
| for i in range(1, len(inverted_residual_setting) - 1): |
| inverted_residual_setting[i][2] = depth_param |
|
|
| if stage_width_list is not None: |
| for i in range(len(inverted_residual_setting)): |
| inverted_residual_setting[i][1] = stage_width_list[i] |
|
|
| ks = val2list(ks, sum([n for _, _, n, _ in inverted_residual_setting]) - 1) |
| _pt = 0 |
|
|
| |
| first_conv = ConvLayer( |
| 3, |
| input_channel, |
| kernel_size=3, |
| stride=2, |
| use_bn=True, |
| act_func="relu6", |
| ops_order="weight_bn_act", |
| ) |
| |
| blocks = [] |
| for t, c, n, s in inverted_residual_setting: |
| output_channel = make_divisible(c * width_mult, MyNetwork.CHANNEL_DIVISIBLE) |
| for i in range(n): |
| if i == 0: |
| stride = s |
| else: |
| stride = 1 |
| if t == 1: |
| kernel_size = 3 |
| else: |
| kernel_size = ks[_pt] |
| _pt += 1 |
| mobile_inverted_conv = MBConvLayer( |
| in_channels=input_channel, |
| out_channels=output_channel, |
| kernel_size=kernel_size, |
| stride=stride, |
| expand_ratio=t, |
| ) |
| if stride == 1: |
| if input_channel == output_channel: |
| shortcut = IdentityLayer(input_channel, input_channel) |
| else: |
| shortcut = ConvLayer(input_channel,output_channel,kernel_size=1,stride=1,bias=False,act_func=None) |
| else: |
| shortcut = None |
| blocks.append(ResidualBlock(mobile_inverted_conv, shortcut)) |
| input_channel = output_channel |
| |
| feature_mix_layer = ConvLayer( |
| input_channel, |
| last_channel, |
| kernel_size=1, |
| use_bn=True, |
| act_func="relu6", |
| ops_order="weight_bn_act", |
| ) |
|
|
| classifier = LinearLayer(last_channel, n_classes, dropout_rate=dropout_rate) |
|
|
| super(MobileNetV2, self).__init__( |
| first_conv, blocks, feature_mix_layer, classifier |
| ) |
|
|
| |
| self.set_bn_param(*bn_param) |
|
|
|
|
|
|
| def proxyless_base_cifar( |
| net_config=None, |
| n_classes=None, |
| bn_param=None, |
| dropout_rate=None, |
| local_path="~/.torch/proxylessnas/", |
| ): |
| assert net_config is not None, "Please input a network config" |
| if "http" in net_config: |
| net_config_path = download_url(net_config, local_path) |
| else: |
| net_config_path = net_config |
| net_config_json = json.load(open(net_config_path, "r")) |
|
|
| if n_classes is not None: |
| net_config_json["classifier"]["out_features"] = n_classes |
| if dropout_rate is not None: |
| net_config_json["classifier"]["dropout_rate"] = dropout_rate |
|
|
| net = ProxylessNASNets_Cifar.build_from_config(net_config_json) |
| if bn_param is not None: |
| net.set_bn_param(*bn_param) |
|
|
| return net |
|
|
|
|
| class ProxylessNASNets_Cifar(MyNetwork): |
| def __init__(self, first_conv, blocks, feature_mix_layer, classifier): |
| super(ProxylessNASNets_Cifar, self).__init__() |
|
|
| self.first_conv = first_conv |
| self.blocks = nn.ModuleList(blocks) |
| self.feature_mix_layer = feature_mix_layer |
| self.global_avg_pool = MyGlobalAvgPool2d(keep_dim=False) |
| self.classifier = classifier |
|
|
| def forward(self, x): |
| x = self.first_conv(x) |
| for block in self.blocks: |
| x = block(x) |
| if self.feature_mix_layer is not None: |
| x = self.feature_mix_layer(x) |
| x = self.global_avg_pool(x) |
| x = self.classifier(x) |
| return x |
|
|
| @property |
| def module_str(self): |
| _str = self.first_conv.module_str + "\n" |
| for block in self.blocks: |
| _str += block.module_str + "\n" |
| _str += self.feature_mix_layer.module_str + "\n" |
| _str += self.global_avg_pool.__repr__() + "\n" |
| _str += self.classifier.module_str |
| return _str |
|
|
| @property |
| def config(self): |
| return { |
| "name": ProxylessNASNets_Cifar.__name__, |
| "bn": self.get_bn_param(), |
| "first_conv": self.first_conv.config, |
| "blocks": [block.config for block in self.blocks], |
| "feature_mix_layer": None |
| if self.feature_mix_layer is None |
| else self.feature_mix_layer.config, |
| "classifier": self.classifier.config, |
| } |
|
|
| @staticmethod |
| def build_from_config(config): |
| first_conv = set_layer_from_config(config["first_conv"]) |
| feature_mix_layer = set_layer_from_config(config["feature_mix_layer"]) |
| classifier = set_layer_from_config(config["classifier"]) |
|
|
| blocks = [] |
| for block_config in config["blocks"]: |
| blocks.append(ResidualBlock.build_from_config(block_config)) |
|
|
| net = ProxylessNASNets_Cifar(first_conv, blocks, feature_mix_layer, classifier) |
| if "bn" in config: |
| net.set_bn_param(**config["bn"]) |
| else: |
| net.set_bn_param(momentum=0.1, eps=1e-3) |
|
|
| return net |
|
|
| def zero_last_gamma(self): |
| for m in self.modules(): |
| if isinstance(m, ResidualBlock): |
| if isinstance(m.conv, MBConvLayer) and isinstance( |
| m.shortcut, IdentityLayer |
| ): |
| m.conv.point_linear.bn.weight.data.zero_() |
|
|
| @property |
| def grouped_block_index(self): |
| info_list = [] |
| block_index_list = [] |
| for i, block in enumerate(self.blocks[1:], 1): |
| if block.shortcut is None and len(block_index_list) > 0: |
| info_list.append(block_index_list) |
| block_index_list = [] |
| block_index_list.append(i) |
| if len(block_index_list) > 0: |
| info_list.append(block_index_list) |
| return info_list |
|
|
| def load_state_dict(self, state_dict, **kwargs): |
| current_state_dict = self.state_dict() |
|
|
| for key in state_dict: |
| if key not in current_state_dict: |
| assert ".mobile_inverted_conv." in key |
| new_key = key.replace(".mobile_inverted_conv.", ".conv.") |
| else: |
| new_key = key |
| current_state_dict[new_key] = state_dict[key] |
| super(ProxylessNASNets_Cifar, self).load_state_dict(current_state_dict) |
|
|
|
|
| class MobileNetV2_Cifar(ProxylessNASNets_Cifar): |
| def __init__( |
| self, |
| n_classes=10, |
| width_mult=1.0, |
| bn_param=(0.1, 1e-3), |
| dropout_rate=0.2, |
| ks=None, |
| expand_ratio=None, |
| depth_param=None, |
| stage_width_list=None, |
| ): |
|
|
| ks = 3 if ks is None else ks |
| expand_ratio = 6 if expand_ratio is None else expand_ratio |
|
|
| input_channel = 32 |
| last_channel = 1280 |
|
|
| input_channel = make_divisible( |
| input_channel * width_mult, MyNetwork.CHANNEL_DIVISIBLE |
| ) |
| last_channel = ( |
| make_divisible(last_channel * width_mult, MyNetwork.CHANNEL_DIVISIBLE) |
| if width_mult > 1.0 |
| else last_channel |
| ) |
|
|
| inverted_residual_setting = [ |
| |
| [1, 16, 1, 1], |
| [expand_ratio, 24, 2, 1], |
| [expand_ratio, 32, 3, 2], |
| [expand_ratio, 64, 4, 2], |
| [expand_ratio, 96, 3, 1], |
| [expand_ratio, 160, 3, 2], |
| [expand_ratio, 320, 1, 1], |
| ] |
|
|
| if depth_param is not None: |
| assert isinstance(depth_param, int) |
| for i in range(1, len(inverted_residual_setting) - 1): |
| inverted_residual_setting[i][2] = depth_param |
|
|
| if stage_width_list is not None: |
| for i in range(len(inverted_residual_setting)): |
| inverted_residual_setting[i][1] = stage_width_list[i] |
|
|
| ks = val2list(ks, sum([n for _, _, n, _ in inverted_residual_setting]) - 1) |
| _pt = 0 |
|
|
| |
| first_conv = ConvLayer( |
| 3, |
| input_channel, |
| kernel_size=3, |
| stride=1, |
| use_bn=True, |
| act_func="relu6", |
| ops_order="weight_bn_act", |
| ) |
| |
| blocks = [] |
| for t, c, n, s in inverted_residual_setting: |
| output_channel = make_divisible(c * width_mult, MyNetwork.CHANNEL_DIVISIBLE) |
| for i in range(n): |
| if i == 0: |
| stride = s |
| else: |
| stride = 1 |
| if t == 1: |
| kernel_size = 3 |
| else: |
| kernel_size = ks[_pt] |
| _pt += 1 |
| mobile_inverted_conv = MBConvLayer( |
| in_channels=input_channel, |
| out_channels=output_channel, |
| kernel_size=kernel_size, |
| stride=stride, |
| expand_ratio=t, |
| ) |
| if stride == 1: |
| if input_channel == output_channel: |
| shortcut = IdentityLayer(input_channel, input_channel) |
| else: |
| shortcut = None |
| else: |
| shortcut = None |
| blocks.append(ResidualBlock(mobile_inverted_conv, shortcut)) |
| input_channel = output_channel |
| |
| feature_mix_layer = ConvLayer( |
| input_channel, |
| last_channel, |
| kernel_size=1, |
| stride=1, |
| use_bn=True, |
| act_func="relu6", |
| ops_order="weight_bn_act", |
| ) |
|
|
| classifier = LinearLayer(last_channel, n_classes, dropout_rate=dropout_rate) |
|
|
| super(MobileNetV2_Cifar, self).__init__( |
| first_conv, blocks, feature_mix_layer, classifier |
| ) |
|
|
| |
| self.set_bn_param(*bn_param) |
|
|
