| |
| |
| |
|
|
| import copy |
| import torch |
| import torch.nn as nn |
| from collections import OrderedDict |
|
|
| from proard.utils.layers import ( |
| MBConvLayer, |
| ConvLayer, |
| IdentityLayer, |
| set_layer_from_config, |
| ) |
| from proard.utils.layers import ResNetBottleneckBlock, LinearLayer |
| from proard.utils import ( |
| MyModule, |
| val2list, |
| get_net_device, |
| build_activation, |
| make_divisible, |
| SEModule, |
| MyNetwork, |
| ) |
| from .dynamic_op import ( |
| DynamicSeparableConv2d, |
| DynamicConv2d, |
| DynamicBatchNorm2d, |
| DynamicSE, |
| DynamicGroupNorm, |
| ) |
| from .dynamic_op import DynamicLinear |
|
|
| __all__ = [ |
| "adjust_bn_according_to_idx", |
| "copy_bn", |
| "DynamicMBConvLayer", |
| "DynamicConvLayer", |
| "DynamicLinearLayer", |
| "DynamicResNetBottleneckBlock", |
| ] |
|
|
|
|
| def adjust_bn_according_to_idx(bn, idx): |
| bn.weight.data = torch.index_select(bn.weight.data, 0, idx) |
| bn.bias.data = torch.index_select(bn.bias.data, 0, idx) |
| if type(bn) in [nn.BatchNorm1d, nn.BatchNorm2d]: |
| bn.running_mean.data = torch.index_select(bn.running_mean.data, 0, idx) |
| bn.running_var.data = torch.index_select(bn.running_var.data, 0, idx) |
|
|
|
|
| def copy_bn(target_bn, src_bn): |
| feature_dim = ( |
| target_bn.num_channels |
| if isinstance(target_bn, nn.GroupNorm) |
| else target_bn.num_features |
| ) |
|
|
| target_bn.weight.data.copy_(src_bn.weight.data[:feature_dim]) |
| target_bn.bias.data.copy_(src_bn.bias.data[:feature_dim]) |
| if type(src_bn) in [nn.BatchNorm1d, nn.BatchNorm2d]: |
| target_bn.running_mean.data.copy_(src_bn.running_mean.data[:feature_dim]) |
| target_bn.running_var.data.copy_(src_bn.running_var.data[:feature_dim]) |
|
|
|
|
| class DynamicLinearLayer(MyModule): |
| def __init__(self, in_features_list, out_features, bias=True, dropout_rate=0): |
| super(DynamicLinearLayer, self).__init__() |
|
|
| self.in_features_list = in_features_list |
| self.out_features = out_features |
| self.bias = bias |
| self.dropout_rate = dropout_rate |
|
|
| if self.dropout_rate > 0: |
| self.dropout = nn.Dropout(self.dropout_rate, inplace=True) |
| else: |
| self.dropout = None |
| self.linear = DynamicLinear( |
| max_in_features=max(self.in_features_list), |
| max_out_features=self.out_features, |
| bias=self.bias, |
| ) |
|
|
| def forward(self, x): |
| if self.dropout is not None: |
| x = self.dropout(x) |
| return self.linear(x) |
|
|
| @property |
| def module_str(self): |
| return "DyLinear(%d, %d)" % (max(self.in_features_list), self.out_features) |
|
|
| @property |
| def config(self): |
| return { |
| "name": DynamicLinear.__name__, |
| "in_features_list": self.in_features_list, |
| "out_features": self.out_features, |
| "bias": self.bias, |
| "dropout_rate": self.dropout_rate, |
| } |
|
|
| @staticmethod |
| def build_from_config(config): |
| return DynamicLinearLayer(**config) |
|
|
| def get_active_subnet(self, in_features, preserve_weight=True): |
| sub_layer = LinearLayer( |
| in_features, self.out_features, self.bias, dropout_rate=self.dropout_rate |
| ) |
| sub_layer = sub_layer.to(get_net_device(self)) |
| if not preserve_weight: |
| return sub_layer |
|
|
| sub_layer.linear.weight.data.copy_( |
| self.linear.get_active_weight(self.out_features, in_features).data |
| ) |
| if self.bias: |
| sub_layer.linear.bias.data.copy_( |
| self.linear.get_active_bias(self.out_features).data |
| ) |
| return sub_layer |
|
|
| def get_active_subnet_config(self, in_features): |
| return { |
| "name": LinearLayer.__name__, |
| "in_features": in_features, |
| "out_features": self.out_features, |
| "bias": self.bias, |
| "dropout_rate": self.dropout_rate, |
| } |
|
|
|
|
| class DynamicMBConvLayer(MyModule): |
| def __init__( |
| self, |
| in_channel_list, |
| out_channel_list, |
| kernel_size_list=3, |
| expand_ratio_list=6, |
| stride=1, |
| act_func="relu6", |
| use_se=False, |
| ): |
| super(DynamicMBConvLayer, self).__init__() |
|
|
| self.in_channel_list = in_channel_list |
| self.out_channel_list = out_channel_list |
|
|
| self.kernel_size_list = val2list(kernel_size_list) |
| self.expand_ratio_list = val2list(expand_ratio_list) |
|
|
| self.stride = stride |
| self.act_func = act_func |
| self.use_se = use_se |
|
|
| |
| max_middle_channel = make_divisible( |
| round(max(self.in_channel_list) * max(self.expand_ratio_list)), |
| MyNetwork.CHANNEL_DIVISIBLE, |
| ) |
| if max(self.expand_ratio_list) == 1: |
| self.inverted_bottleneck = None |
| else: |
| self.inverted_bottleneck = nn.Sequential( |
| OrderedDict( |
| [ |
| ( |
| "conv", |
| DynamicConv2d( |
| max(self.in_channel_list), max_middle_channel |
| ), |
| ), |
| ("bn", DynamicBatchNorm2d(max_middle_channel)), |
| ("act", build_activation(self.act_func)), |
| ] |
| ) |
| ) |
|
|
| self.depth_conv = nn.Sequential( |
| OrderedDict( |
| [ |
| ( |
| "conv", |
| DynamicSeparableConv2d( |
| max_middle_channel, self.kernel_size_list, self.stride |
| ), |
| ), |
| ("bn", DynamicBatchNorm2d(max_middle_channel)), |
| ("act", build_activation(self.act_func)), |
| ] |
| ) |
| ) |
| if self.use_se: |
| self.depth_conv.add_module("se", DynamicSE(max_middle_channel)) |
|
|
| self.point_linear = nn.Sequential( |
| OrderedDict( |
| [ |
| ( |
| "conv", |
| DynamicConv2d(max_middle_channel, max(self.out_channel_list)), |
| ), |
| ("bn", DynamicBatchNorm2d(max(self.out_channel_list))), |
| ] |
| ) |
| ) |
|
|
| self.active_kernel_size = max(self.kernel_size_list) |
| self.active_expand_ratio = max(self.expand_ratio_list) |
| self.active_out_channel = max(self.out_channel_list) |
|
|
| def forward(self, x): |
| in_channel = x.size(1) |
|
|
| if self.inverted_bottleneck is not None: |
| self.inverted_bottleneck.conv.active_out_channel = make_divisible( |
| round(in_channel * self.active_expand_ratio), |
| MyNetwork.CHANNEL_DIVISIBLE, |
| ) |
|
|
| self.depth_conv.conv.active_kernel_size = self.active_kernel_size |
| self.point_linear.conv.active_out_channel = self.active_out_channel |
|
|
| if self.inverted_bottleneck is not None: |
| x = self.inverted_bottleneck(x) |
| x = self.depth_conv(x) |
| x = self.point_linear(x) |
| return x |
|
|
| @property |
| def module_str(self): |
| if self.use_se: |
| return "SE(O%d, E%.1f, K%d)" % ( |
| self.active_out_channel, |
| self.active_expand_ratio, |
| self.active_kernel_size, |
| ) |
| else: |
| return "(O%d, E%.1f, K%d)" % ( |
| self.active_out_channel, |
| self.active_expand_ratio, |
| self.active_kernel_size, |
| ) |
|
|
| @property |
| def config(self): |
| return { |
| "name": DynamicMBConvLayer.__name__, |
| "in_channel_list": self.in_channel_list, |
| "out_channel_list": self.out_channel_list, |
| "kernel_size_list": self.kernel_size_list, |
| "expand_ratio_list": self.expand_ratio_list, |
| "stride": self.stride, |
| "act_func": self.act_func, |
| "use_se": self.use_se, |
| } |
|
|
| @staticmethod |
| def build_from_config(config): |
| return DynamicMBConvLayer(**config) |
|
|
| |
|
|
| @property |
| def in_channels(self): |
| return max(self.in_channel_list) |
|
|
| @property |
| def out_channels(self): |
| return max(self.out_channel_list) |
|
|
| def active_middle_channel(self, in_channel): |
| return make_divisible( |
| round(in_channel * self.active_expand_ratio), MyNetwork.CHANNEL_DIVISIBLE |
| ) |
|
|
| |
|
|
| def get_active_subnet(self, in_channel, preserve_weight=True): |
| |
| sub_layer = set_layer_from_config(self.get_active_subnet_config(in_channel)) |
| sub_layer = sub_layer.to(get_net_device(self)) |
| if not preserve_weight: |
| return sub_layer |
|
|
| middle_channel = self.active_middle_channel(in_channel) |
| |
| if sub_layer.inverted_bottleneck is not None: |
| sub_layer.inverted_bottleneck.conv.weight.data.copy_( |
| self.inverted_bottleneck.conv.get_active_filter( |
| middle_channel, in_channel |
| ).data, |
| ) |
| copy_bn(sub_layer.inverted_bottleneck.bn, self.inverted_bottleneck.bn.bn) |
|
|
| sub_layer.depth_conv.conv.weight.data.copy_( |
| self.depth_conv.conv.get_active_filter( |
| middle_channel, self.active_kernel_size |
| ).data |
| ) |
| copy_bn(sub_layer.depth_conv.bn, self.depth_conv.bn.bn) |
|
|
| if self.use_se: |
| se_mid = make_divisible( |
| middle_channel // SEModule.REDUCTION, |
| divisor=MyNetwork.CHANNEL_DIVISIBLE, |
| ) |
| sub_layer.depth_conv.se.fc.reduce.weight.data.copy_( |
| self.depth_conv.se.get_active_reduce_weight(se_mid, middle_channel).data |
| ) |
| sub_layer.depth_conv.se.fc.reduce.bias.data.copy_( |
| self.depth_conv.se.get_active_reduce_bias(se_mid).data |
| ) |
|
|
| sub_layer.depth_conv.se.fc.expand.weight.data.copy_( |
| self.depth_conv.se.get_active_expand_weight(se_mid, middle_channel).data |
| ) |
| sub_layer.depth_conv.se.fc.expand.bias.data.copy_( |
| self.depth_conv.se.get_active_expand_bias(middle_channel).data |
| ) |
|
|
| sub_layer.point_linear.conv.weight.data.copy_( |
| self.point_linear.conv.get_active_filter( |
| self.active_out_channel, middle_channel |
| ).data |
| ) |
| copy_bn(sub_layer.point_linear.bn, self.point_linear.bn.bn) |
|
|
| return sub_layer |
|
|
| def get_active_subnet_config(self, in_channel): |
| return { |
| "name": MBConvLayer.__name__, |
| "in_channels": in_channel, |
| "out_channels": self.active_out_channel, |
| "kernel_size": self.active_kernel_size, |
| "stride": self.stride, |
| "expand_ratio": self.active_expand_ratio, |
| "mid_channels": self.active_middle_channel(in_channel), |
| "act_func": self.act_func, |
| "use_se": self.use_se, |
| } |
|
|
| def re_organize_middle_weights(self, expand_ratio_stage=0): |
| importance = torch.sum( |
| torch.abs(self.point_linear.conv.conv.weight.data), dim=(0, 2, 3) |
| ) |
| if isinstance(self.depth_conv.bn, DynamicGroupNorm): |
| channel_per_group = self.depth_conv.bn.channel_per_group |
| importance_chunks = torch.split(importance, channel_per_group) |
| for chunk in importance_chunks: |
| chunk.data.fill_(torch.mean(chunk)) |
| importance = torch.cat(importance_chunks, dim=0) |
| if expand_ratio_stage > 0: |
| sorted_expand_list = copy.deepcopy(self.expand_ratio_list) |
| sorted_expand_list.sort(reverse=True) |
| target_width_list = [ |
| make_divisible( |
| round(max(self.in_channel_list) * expand), |
| MyNetwork.CHANNEL_DIVISIBLE, |
| ) |
| for expand in sorted_expand_list |
| ] |
|
|
| right = len(importance) |
| base = -len(target_width_list) * 1e5 |
| for i in range(expand_ratio_stage + 1): |
| left = target_width_list[i] |
| importance[left:right] += base |
| base += 1e5 |
| right = left |
|
|
| sorted_importance, sorted_idx = torch.sort(importance, dim=0, descending=True) |
| self.point_linear.conv.conv.weight.data = torch.index_select( |
| self.point_linear.conv.conv.weight.data, 1, sorted_idx |
| ) |
|
|
| adjust_bn_according_to_idx(self.depth_conv.bn.bn, sorted_idx) |
| self.depth_conv.conv.conv.weight.data = torch.index_select( |
| self.depth_conv.conv.conv.weight.data, 0, sorted_idx |
| ) |
|
|
| if self.use_se: |
| |
| se_expand = self.depth_conv.se.fc.expand |
| se_expand.weight.data = torch.index_select( |
| se_expand.weight.data, 0, sorted_idx |
| ) |
| se_expand.bias.data = torch.index_select(se_expand.bias.data, 0, sorted_idx) |
| |
| se_reduce = self.depth_conv.se.fc.reduce |
| se_reduce.weight.data = torch.index_select( |
| se_reduce.weight.data, 1, sorted_idx |
| ) |
| |
| se_importance = torch.sum(torch.abs(se_expand.weight.data), dim=(0, 2, 3)) |
| se_importance, se_idx = torch.sort(se_importance, dim=0, descending=True) |
|
|
| se_expand.weight.data = torch.index_select(se_expand.weight.data, 1, se_idx) |
| se_reduce.weight.data = torch.index_select(se_reduce.weight.data, 0, se_idx) |
| se_reduce.bias.data = torch.index_select(se_reduce.bias.data, 0, se_idx) |
|
|
| if self.inverted_bottleneck is not None: |
| adjust_bn_according_to_idx(self.inverted_bottleneck.bn.bn, sorted_idx) |
| self.inverted_bottleneck.conv.conv.weight.data = torch.index_select( |
| self.inverted_bottleneck.conv.conv.weight.data, 0, sorted_idx |
| ) |
| return None |
| else: |
| return sorted_idx |
|
|
|
|
| class DynamicConvLayer(MyModule): |
| def __init__( |
| self, |
| in_channel_list, |
| out_channel_list, |
| kernel_size=3, |
| stride=1, |
| dilation=1, |
| use_bn=True, |
| act_func="relu6", |
| ): |
| super(DynamicConvLayer, self).__init__() |
|
|
| self.in_channel_list = in_channel_list |
| self.out_channel_list = out_channel_list |
| self.kernel_size = kernel_size |
| self.stride = stride |
| self.dilation = dilation |
| self.use_bn = use_bn |
| self.act_func = act_func |
|
|
| self.conv = DynamicConv2d( |
| max_in_channels=max(self.in_channel_list), |
| max_out_channels=max(self.out_channel_list), |
| kernel_size=self.kernel_size, |
| stride=self.stride, |
| dilation=self.dilation, |
| ) |
| if self.use_bn: |
| self.bn = DynamicBatchNorm2d(max(self.out_channel_list)) |
| self.act = build_activation(self.act_func) |
|
|
| self.active_out_channel = max(self.out_channel_list) |
|
|
| def forward(self, x): |
| self.conv.active_out_channel = self.active_out_channel |
|
|
| x = self.conv(x) |
| if self.use_bn: |
| x = self.bn(x) |
| x = self.act(x) |
| return x |
|
|
| @property |
| def module_str(self): |
| return "DyConv(O%d, K%d, S%d)" % ( |
| self.active_out_channel, |
| self.kernel_size, |
| self.stride, |
| ) |
|
|
| @property |
| def config(self): |
| return { |
| "name": DynamicConvLayer.__name__, |
| "in_channel_list": self.in_channel_list, |
| "out_channel_list": self.out_channel_list, |
| "kernel_size": self.kernel_size, |
| "stride": self.stride, |
| "dilation": self.dilation, |
| "use_bn": self.use_bn, |
| "act_func": self.act_func, |
| } |
|
|
| @staticmethod |
| def build_from_config(config): |
| return DynamicConvLayer(**config) |
|
|
| |
|
|
| @property |
| def in_channels(self): |
| return max(self.in_channel_list) |
|
|
| @property |
| def out_channels(self): |
| return max(self.out_channel_list) |
|
|
| |
|
|
| def get_active_subnet(self, in_channel, preserve_weight=True): |
| sub_layer = set_layer_from_config(self.get_active_subnet_config(in_channel)) |
| sub_layer = sub_layer.to(get_net_device(self)) |
|
|
| if not preserve_weight: |
| return sub_layer |
|
|
| sub_layer.conv.weight.data.copy_( |
| self.conv.get_active_filter(self.active_out_channel, in_channel).data |
| ) |
| if self.use_bn: |
| copy_bn(sub_layer.bn, self.bn.bn) |
|
|
| return sub_layer |
|
|
| def get_active_subnet_config(self, in_channel): |
| return { |
| "name": ConvLayer.__name__, |
| "in_channels": in_channel, |
| "out_channels": self.active_out_channel, |
| "kernel_size": self.kernel_size, |
| "stride": self.stride, |
| "dilation": self.dilation, |
| "use_bn": self.use_bn, |
| "act_func": self.act_func, |
| } |
|
|
|
|
| class DynamicResNetBottleneckBlock(MyModule): |
| def __init__( |
| self, |
| in_channel_list, |
| out_channel_list, |
| expand_ratio_list=0.25, |
| kernel_size=3, |
| stride=1, |
| act_func="relu", |
| downsample_mode="avgpool_conv", |
| ): |
| super(DynamicResNetBottleneckBlock, self).__init__() |
|
|
| self.in_channel_list = in_channel_list |
| self.out_channel_list = out_channel_list |
| self.expand_ratio_list = val2list(expand_ratio_list) |
|
|
| self.kernel_size = kernel_size |
| self.stride = stride |
| self.act_func = act_func |
| self.downsample_mode = downsample_mode |
|
|
| |
| max_middle_channel = make_divisible( |
| round(max(self.out_channel_list) * max(self.expand_ratio_list)), |
| MyNetwork.CHANNEL_DIVISIBLE, |
| ) |
|
|
| self.conv1 = nn.Sequential( |
| OrderedDict( |
| [ |
| ( |
| "conv", |
| DynamicConv2d(max(self.in_channel_list), max_middle_channel), |
| ), |
| ("bn", DynamicBatchNorm2d(max_middle_channel)), |
| ("act", build_activation(self.act_func, inplace=True)), |
| ] |
| ) |
| ) |
|
|
| self.conv2 = nn.Sequential( |
| OrderedDict( |
| [ |
| ( |
| "conv", |
| DynamicConv2d( |
| max_middle_channel, max_middle_channel, kernel_size, stride |
| ), |
| ), |
| ("bn", DynamicBatchNorm2d(max_middle_channel)), |
| ("act", build_activation(self.act_func, inplace=True)), |
| ] |
| ) |
| ) |
|
|
| self.conv3 = nn.Sequential( |
| OrderedDict( |
| [ |
| ( |
| "conv", |
| DynamicConv2d(max_middle_channel, max(self.out_channel_list)), |
| ), |
| ("bn", DynamicBatchNorm2d(max(self.out_channel_list))), |
| ] |
| ) |
| ) |
|
|
| if self.stride == 1 and self.in_channel_list == self.out_channel_list: |
| self.downsample = IdentityLayer( |
| max(self.in_channel_list), max(self.out_channel_list) |
| ) |
| elif self.downsample_mode == "conv": |
| self.downsample = nn.Sequential( |
| OrderedDict( |
| [ |
| ( |
| "conv", |
| DynamicConv2d( |
| max(self.in_channel_list), |
| max(self.out_channel_list), |
| stride=stride, |
| ), |
| ), |
| ("bn", DynamicBatchNorm2d(max(self.out_channel_list))), |
| ] |
| ) |
| ) |
| elif self.downsample_mode == "avgpool_conv": |
| self.downsample = nn.Sequential( |
| OrderedDict( |
| [ |
| ( |
| "avg_pool", |
| nn.AvgPool2d( |
| kernel_size=stride, |
| stride=stride, |
| padding=0, |
| ceil_mode=True, |
| ), |
| ), |
| ( |
| "conv", |
| DynamicConv2d( |
| max(self.in_channel_list), max(self.out_channel_list) |
| ), |
| ), |
| ("bn", DynamicBatchNorm2d(max(self.out_channel_list))), |
| ] |
| ) |
| ) |
| else: |
| raise NotImplementedError |
|
|
| self.final_act = build_activation(self.act_func, inplace=True) |
|
|
| self.active_expand_ratio = max(self.expand_ratio_list) |
| self.active_out_channel = max(self.out_channel_list) |
|
|
| def forward(self, x): |
| feature_dim = self.active_middle_channels |
|
|
| self.conv1.conv.active_out_channel = feature_dim |
| self.conv2.conv.active_out_channel = feature_dim |
| self.conv3.conv.active_out_channel = self.active_out_channel |
| if not isinstance(self.downsample, IdentityLayer): |
| self.downsample.conv.active_out_channel = self.active_out_channel |
|
|
| residual = self.downsample(x) |
|
|
| x = self.conv1(x) |
| x = self.conv2(x) |
| x = self.conv3(x) |
|
|
| x = x + residual |
| x = self.final_act(x) |
| return x |
|
|
| @property |
| def module_str(self): |
| return "(%s, %s)" % ( |
| "%dx%d_BottleneckConv_in->%d->%d_S%d" |
| % ( |
| self.kernel_size, |
| self.kernel_size, |
| self.active_middle_channels, |
| self.active_out_channel, |
| self.stride, |
| ), |
| "Identity" |
| if isinstance(self.downsample, IdentityLayer) |
| else self.downsample_mode, |
| ) |
|
|
| @property |
| def config(self): |
| return { |
| "name": DynamicResNetBottleneckBlock.__name__, |
| "in_channel_list": self.in_channel_list, |
| "out_channel_list": self.out_channel_list, |
| "expand_ratio_list": self.expand_ratio_list, |
| "kernel_size": self.kernel_size, |
| "stride": self.stride, |
| "act_func": self.act_func, |
| "downsample_mode": self.downsample_mode, |
| } |
|
|
| @staticmethod |
| def build_from_config(config): |
| return DynamicResNetBottleneckBlock(**config) |
|
|
| |
|
|
| @property |
| def in_channels(self): |
| return max(self.in_channel_list) |
|
|
| @property |
| def out_channels(self): |
| return max(self.out_channel_list) |
|
|
| @property |
| def active_middle_channels(self): |
| feature_dim = round(self.active_out_channel * self.active_expand_ratio) |
| feature_dim = make_divisible(feature_dim, MyNetwork.CHANNEL_DIVISIBLE) |
| return feature_dim |
|
|
| |
|
|
| def get_active_subnet(self, in_channel, preserve_weight=True): |
| |
| sub_layer = set_layer_from_config(self.get_active_subnet_config(in_channel)) |
| sub_layer = sub_layer.to(get_net_device(self)) |
| if not preserve_weight: |
| return sub_layer |
|
|
| |
| sub_layer.conv1.conv.weight.data.copy_( |
| self.conv1.conv.get_active_filter( |
| self.active_middle_channels, in_channel |
| ).data |
| ) |
| copy_bn(sub_layer.conv1.bn, self.conv1.bn.bn) |
|
|
| sub_layer.conv2.conv.weight.data.copy_( |
| self.conv2.conv.get_active_filter( |
| self.active_middle_channels, self.active_middle_channels |
| ).data |
| ) |
| copy_bn(sub_layer.conv2.bn, self.conv2.bn.bn) |
|
|
| sub_layer.conv3.conv.weight.data.copy_( |
| self.conv3.conv.get_active_filter( |
| self.active_out_channel, self.active_middle_channels |
| ).data |
| ) |
| copy_bn(sub_layer.conv3.bn, self.conv3.bn.bn) |
|
|
| if not isinstance(self.downsample, IdentityLayer): |
| sub_layer.downsample.conv.weight.data.copy_( |
| self.downsample.conv.get_active_filter( |
| self.active_out_channel, in_channel |
| ).data |
| ) |
| copy_bn(sub_layer.downsample.bn, self.downsample.bn.bn) |
|
|
| return sub_layer |
|
|
| def get_active_subnet_config(self, in_channel): |
| return { |
| "name": ResNetBottleneckBlock.__name__, |
| "in_channels": in_channel, |
| "out_channels": self.active_out_channel, |
| "kernel_size": self.kernel_size, |
| "stride": self.stride, |
| "expand_ratio": self.active_expand_ratio, |
| "mid_channels": self.active_middle_channels, |
| "act_func": self.act_func, |
| "groups": 1, |
| "downsample_mode": self.downsample_mode, |
| } |
|
|
| def re_organize_middle_weights(self, expand_ratio_stage=0): |
| |
| importance = torch.sum( |
| torch.abs(self.conv3.conv.conv.weight.data), dim=(0, 2, 3) |
| ) |
| if isinstance(self.conv2.bn, DynamicGroupNorm): |
| channel_per_group = self.conv2.bn.channel_per_group |
| importance_chunks = torch.split(importance, channel_per_group) |
| for chunk in importance_chunks: |
| chunk.data.fill_(torch.mean(chunk)) |
| importance = torch.cat(importance_chunks, dim=0) |
| if expand_ratio_stage > 0: |
| sorted_expand_list = copy.deepcopy(self.expand_ratio_list) |
| sorted_expand_list.sort(reverse=True) |
| target_width_list = [ |
| make_divisible( |
| round(max(self.out_channel_list) * expand), |
| MyNetwork.CHANNEL_DIVISIBLE, |
| ) |
| for expand in sorted_expand_list |
| ] |
| right = len(importance) |
| base = -len(target_width_list) * 1e5 |
| for i in range(expand_ratio_stage + 1): |
| left = target_width_list[i] |
| importance[left:right] += base |
| base += 1e5 |
| right = left |
|
|
| sorted_importance, sorted_idx = torch.sort(importance, dim=0, descending=True) |
| self.conv3.conv.conv.weight.data = torch.index_select( |
| self.conv3.conv.conv.weight.data, 1, sorted_idx |
| ) |
| adjust_bn_according_to_idx(self.conv2.bn.bn, sorted_idx) |
| self.conv2.conv.conv.weight.data = torch.index_select( |
| self.conv2.conv.conv.weight.data, 0, sorted_idx |
| ) |
|
|
| |
| importance = torch.sum( |
| torch.abs(self.conv2.conv.conv.weight.data), dim=(0, 2, 3) |
| ) |
| if isinstance(self.conv1.bn, DynamicGroupNorm): |
| channel_per_group = self.conv1.bn.channel_per_group |
| importance_chunks = torch.split(importance, channel_per_group) |
| for chunk in importance_chunks: |
| chunk.data.fill_(torch.mean(chunk)) |
| importance = torch.cat(importance_chunks, dim=0) |
| if expand_ratio_stage > 0: |
| sorted_expand_list = copy.deepcopy(self.expand_ratio_list) |
| sorted_expand_list.sort(reverse=True) |
| target_width_list = [ |
| make_divisible( |
| round(max(self.out_channel_list) * expand), |
| MyNetwork.CHANNEL_DIVISIBLE, |
| ) |
| for expand in sorted_expand_list |
| ] |
| right = len(importance) |
| base = -len(target_width_list) * 1e5 |
| for i in range(expand_ratio_stage + 1): |
| left = target_width_list[i] |
| importance[left:right] += base |
| base += 1e5 |
| right = left |
| sorted_importance, sorted_idx = torch.sort(importance, dim=0, descending=True) |
|
|
| self.conv2.conv.conv.weight.data = torch.index_select( |
| self.conv2.conv.conv.weight.data, 1, sorted_idx |
| ) |
| adjust_bn_according_to_idx(self.conv1.bn.bn, sorted_idx) |
| self.conv1.conv.conv.weight.data = torch.index_select( |
| self.conv1.conv.conv.weight.data, 0, sorted_idx |
| ) |
|
|
| return None |
|
|
