• 使用NNI对DLASeg剪枝的失败记录

    本文希望对CenterNet算法的Backbone暨DLASeg进行剪枝。

    剪枝试验涉及3个文件，分别为：

    DCN可变性卷积dcn_v2.py，因为DLASeg依赖DCN。

    1. #!/usr/bin/env python
    2. from __future__ import absolute_import
    3. from __future__ import print_function
    4. from __future__ import division
    5.  
    6. import math
    7. import torch
    8. from torch import nn
    9. from torch.autograd import Function
    10. from torch.nn.modules.utils import _pair
    11. from torch.autograd.function import once_differentiable
    12.  
    13. import _ext as _backend
    14.  
    15.  
    16. class _DCNv2(Function):
    17.     @staticmethod
    18.     def forward(ctx, input, offset, mask, weight, bias,
    19.                 stride, padding, dilation, deformable_groups):
    20.         ctx.stride = _pair(stride)
    21.         ctx.padding = _pair(padding)
    22.         ctx.dilation = _pair(dilation)
    23.         ctx.kernel_size = _pair(weight.shape[2:4])
    24.         ctx.deformable_groups = deformable_groups
    25.         output = _backend.dcn_v2_forward(input, weight, bias,
    26.                                          offset, mask,
    27.                                          ctx.kernel_size[0], ctx.kernel_size[1],
    28.                                          ctx.stride[0], ctx.stride[1],
    29.                                          ctx.padding[0], ctx.padding[1],
    30.                                          ctx.dilation[0], ctx.dilation[1],
    31.                                          ctx.deformable_groups)
    32.         ctx.save_for_backward(input, offset, mask, weight, bias)
    33.         return output
    34.  
    35.     @staticmethod
    36.     @once_differentiable
    37.     def backward(ctx, grad_output):
    38.         input, offset, mask, weight, bias = ctx.saved_tensors
    39.         grad_input, grad_offset, grad_mask, grad_weight, grad_bias = \
    40.             _backend.dcn_v2_backward(input, weight,
    41.                                      bias,
    42.                                      offset, mask,
    43.                                      grad_output,
    44.                                      ctx.kernel_size[0], ctx.kernel_size[1],
    45.                                      ctx.stride[0], ctx.stride[1],
    46.                                      ctx.padding[0], ctx.padding[1],
    47.                                      ctx.dilation[0], ctx.dilation[1],
    48.                                      ctx.deformable_groups)
    49.  
    50.         return grad_input, grad_offset, grad_mask, grad_weight, grad_bias,\
    51.             None, None, None, None,
    52.  
    53.  
    54. dcn_v2_conv = _DCNv2.apply
    55.  
    56.  
    57. class DCNv2(nn.Module):
    58.  
    59.     def __init__(self, in_channels, out_channels,
    60.                  kernel_size, stride, padding, dilation=1, deformable_groups=1):
    61.         super(DCNv2, self).__init__()
    62.         self.in_channels = in_channels
    63.         self.out_channels = out_channels
    64.         self.kernel_size = _pair(kernel_size)
    65.         self.stride = _pair(stride)
    66.         self.padding = _pair(padding)
    67.         self.dilation = _pair(dilation)
    68.         self.deformable_groups = deformable_groups
    69.  
    70.         self.weight = nn.Parameter(torch.Tensor(
    71.             out_channels, in_channels, *self.kernel_size))
    72.         self.bias = nn.Parameter(torch.Tensor(out_channels))
    73.         self.reset_parameters()
    74.  
    75.     def reset_parameters(self):
    76.         n = self.in_channels
    77.         for k in self.kernel_size:
    78.             n *= k
    79.         stdv = 1. / math.sqrt(n)
    80.         self.weight.data.uniform_(-stdv, stdv)
    81.         self.bias.data.zero_()
    82.  
    83.     def forward(self, input, offset, mask):
    84.         assert 2 * self.deformable_groups * self.kernel_size[0] * self.kernel_size[1] == \
    85.             offset.shape[1]
    86.         assert self.deformable_groups * self.kernel_size[0] * self.kernel_size[1] == \
    87.             mask.shape[1]
    88.         return dcn_v2_conv(input, offset, mask,
    89.                            self.weight,
    90.                            self.bias,
    91.                            self.stride,
    92.                            self.padding,
    93.                            self.dilation,
    94.                            self.deformable_groups)
    95.  
    96.  
    97. class DCN(DCNv2):
    98.  
    99.     def __init__(self, in_channels, out_channels,
    100.                  kernel_size, stride, padding,
    101.                  dilation=1, deformable_groups=1):
    102.         super(DCN, self).__init__(in_channels, out_channels,
    103.                                   kernel_size, stride, padding, dilation, deformable_groups)
    104.  
    105.         channels_ = self.deformable_groups * 3 * self.kernel_size[0] * self.kernel_size[1]
    106.         self.conv_offset_mask = nn.Conv2d(self.in_channels,
    107.                                           channels_,
    108.                                           kernel_size=self.kernel_size,
    109.                                           stride=self.stride,
    110.                                           padding=self.padding,
    111.                                           bias=True)
    112.         self.init_offset()
    113.  
    114.     def init_offset(self):
    115.         self.conv_offset_mask.weight.data.zero_()
    116.         self.conv_offset_mask.bias.data.zero_()
    117.  
    118.     def forward(self, input):
    119.         out = self.conv_offset_mask(input)
    120.         o1, o2, mask = torch.chunk(out, 3, dim=1)
    121.         offset = torch.cat((o1, o2), dim=1)
    122.         mask = torch.sigmoid(mask)
    123.         return dcn_v2_conv(input, offset, mask,
    124.                            self.weight, self.bias,
    125.                            self.stride,
    126.                            self.padding,
    127.                            self.dilation,
    128.                            self.deformable_groups)
    129.  
    130.  
    131.  
    132. class _DCNv2Pooling(Function):
    133.     @staticmethod
    134.     def forward(ctx, input, rois, offset,
    135.                 spatial_scale,
    136.                 pooled_size,
    137.                 output_dim,
    138.                 no_trans,
    139.                 group_size=1,
    140.                 part_size=None,
    141.                 sample_per_part=4,
    142.                 trans_std=.0):
    143.         ctx.spatial_scale = spatial_scale
    144.         ctx.no_trans = int(no_trans)
    145.         ctx.output_dim = output_dim
    146.         ctx.group_size = group_size
    147.         ctx.pooled_size = pooled_size
    148.         ctx.part_size = pooled_size if part_size is None else part_size
    149.         ctx.sample_per_part = sample_per_part
    150.         ctx.trans_std = trans_std
    151.  
    152.         output, output_count = \
    153.             _backend.dcn_v2_psroi_pooling_forward(input, rois, offset,
    154.                                                   ctx.no_trans, ctx.spatial_scale,
    155.                                                   ctx.output_dim, ctx.group_size,
    156.                                                   ctx.pooled_size, ctx.part_size,
    157.                                                   ctx.sample_per_part, ctx.trans_std)
    158.         ctx.save_for_backward(input, rois, offset, output_count)
    159.         return output
    160.  
    161.     @staticmethod
    162.     @once_differentiable
    163.     def backward(ctx, grad_output):
    164.         input, rois, offset, output_count = ctx.saved_tensors
    165.         grad_input, grad_offset = \
    166.             _backend.dcn_v2_psroi_pooling_backward(grad_output,
    167.                                                    input,
    168.                                                    rois,
    169.                                                    offset,
    170.                                                    output_count,
    171.                                                    ctx.no_trans,
    172.                                                    ctx.spatial_scale,
    173.                                                    ctx.output_dim,
    174.                                                    ctx.group_size,
    175.                                                    ctx.pooled_size,
    176.                                                    ctx.part_size,
    177.                                                    ctx.sample_per_part,
    178.                                                    ctx.trans_std)
    179.  
    180.         return grad_input, None, grad_offset, \
    181.             None, None, None, None, None, None, None, None
    182.  
    183.  
    184. dcn_v2_pooling = _DCNv2Pooling.apply
    185.  
    186.  
    187. class DCNv2Pooling(nn.Module):
    188.  
    189.     def __init__(self,
    190.                  spatial_scale,
    191.                  pooled_size,
    192.                  output_dim,
    193.                  no_trans,
    194.                  group_size=1,
    195.                  part_size=None,
    196.                  sample_per_part=4,
    197.                  trans_std=.0):
    198.         super(DCNv2Pooling, self).__init__()
    199.         self.spatial_scale = spatial_scale
    200.         self.pooled_size = pooled_size
    201.         self.output_dim = output_dim
    202.         self.no_trans = no_trans
    203.         self.group_size = group_size
    204.         self.part_size = pooled_size if part_size is None else part_size
    205.         self.sample_per_part = sample_per_part
    206.         self.trans_std = trans_std
    207.  
    208.     def forward(self, input, rois, offset):
    209.         assert input.shape[1] == self.output_dim
    210.         if self.no_trans:
    211.             offset = input.new()
    212.         return dcn_v2_pooling(input, rois, offset,
    213.                               self.spatial_scale,
    214.                               self.pooled_size,
    215.                               self.output_dim,
    216.                               self.no_trans,
    217.                               self.group_size,
    218.                               self.part_size,
    219.                               self.sample_per_part,
    220.                               self.trans_std)
    221.  
    222.  
    223. class DCNPooling(DCNv2Pooling):
    224.  
    225.     def __init__(self,
    226.                  spatial_scale,
    227.                  pooled_size,
    228.                  output_dim,
    229.                  no_trans,
    230.                  group_size=1,
    231.                  part_size=None,
    232.                  sample_per_part=4,
    233.                  trans_std=.0,
    234.                  deform_fc_dim=1024):
    235.         super(DCNPooling, self).__init__(spatial_scale,
    236.                                          pooled_size,
    237.                                          output_dim,
    238.                                          no_trans,
    239.                                          group_size,
    240.                                          part_size,
    241.                                          sample_per_part,
    242.                                          trans_std)
    243.  
    244.         self.deform_fc_dim = deform_fc_dim
    245.  
    246.         if not no_trans:
    247.             self.offset_mask_fc = nn.Sequential(
    248.                 nn.Linear(self.pooled_size * self.pooled_size *
    249.                           self.output_dim, self.deform_fc_dim),
    250.                 nn.ReLU(inplace=True),
    251.                 nn.Linear(self.deform_fc_dim, self.deform_fc_dim),
    252.                 nn.ReLU(inplace=True),
    253.                 nn.Linear(self.deform_fc_dim, self.pooled_size *
    254.                           self.pooled_size * 3)
    255.             )
    256.             self.offset_mask_fc[4].weight.data.zero_()
    257.             self.offset_mask_fc[4].bias.data.zero_()
    258.  
    259.     def forward(self, input, rois):
    260.         offset = input.new()
    261.  
    262.         if not self.no_trans:
    263.  
    264.             # do roi_align first
    265.             n = rois.shape[0]
    266.             roi = dcn_v2_pooling(input, rois, offset,
    267.                                  self.spatial_scale,
    268.                                  self.pooled_size,
    269.                                  self.output_dim,
    270.                                  True,  # no trans
    271.                                  self.group_size,
    272.                                  self.part_size,
    273.                                  self.sample_per_part,
    274.                                  self.trans_std)
    275.  
    276.             # build mask and offset
    277.             offset_mask = self.offset_mask_fc(roi.view(n, -1))
    278.             offset_mask = offset_mask.view(
    279.                 n, 3, self.pooled_size, self.pooled_size)
    280.             o1, o2, mask = torch.chunk(offset_mask, 3, dim=1)
    281.             offset = torch.cat((o1, o2), dim=1)
    282.             mask = torch.sigmoid(mask)
    283.  
    284.             # do pooling with offset and mask
    285.             return dcn_v2_pooling(input, rois, offset,
    286.                                   self.spatial_scale,
    287.                                   self.pooled_size,
    288.                                   self.output_dim,
    289.                                   self.no_trans,
    290.                                   self.group_size,
    291.                                   self.part_size,
    292.                                   self.sample_per_part,
    293.                                   self.trans_std) * mask
    294.         # only roi_align
    295.         return dcn_v2_pooling(input, rois, offset,
    296.                               self.spatial_scale,
    297.                               self.pooled_size,
    298.                               self.output_dim,
    299.                               self.no_trans,
    300.                               self.group_size,
    301.                               self.part_size,
    302.                               self.sample_per_part,
    303.                               self.trans_std)

    pose_dla_dcn.py文件：

    1. from __future__ import absolute_import
    2. from __future__ import division
    3. from __future__ import print_function
    4.  
    5. import os
    6. import math
    7. import logging
    8. import numpy as np
    9. from os.path import join
    10.  
    11. import torch
    12. from torch import nn
    13. import torch.nn.functional as F
    14. import torch.utils.model_zoo as model_zoo
    15.  
    16. from dcn_v2 import DCN
    17.  
    18. BN_MOMENTUM = 0.1
    19. logger = logging.getLogger(__name__)
    20.  
    21. def get_model_url(data='imagenet', name='dla34', hash='ba72cf86'):
    22.     return join('http://dl.yf.io/dla/models', data, '{}-{}.pth'.format(name, hash))
    23.  
    24.  
    25. def conv3x3(in_planes, out_planes, stride=1):
    26.     "3x3 convolution with padding"
    27.     return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
    28.                      padding=1, bias=False)
    29.  
    30.  
    31. class BasicBlock(nn.Module):
    32.     def __init__(self, inplanes, planes, stride=1, dilation=1):
    33.         super(BasicBlock, self).__init__()
    34.         self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3,
    35.                                stride=stride, padding=dilation,
    36.                                bias=False, dilation=dilation)
    37.         self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
    38.         self.relu = nn.ReLU(inplace=True)
    39.         self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,
    40.                                stride=1, padding=dilation,
    41.                                bias=False, dilation=dilation)
    42.         self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
    43.         self.stride = stride
    44.  
    45.     def forward(self, x, residual=None):
    46.         if residual is None:
    47.             residual = x
    48.  
    49.         out = self.conv1(x)
    50.         out = self.bn1(out)
    51.         out = self.relu(out)
    52.  
    53.         out = self.conv2(out)
    54.         out = self.bn2(out)
    55.  
    56.         out += residual
    57.         out = self.relu(out)
    58.  
    59.         return out
    60.  
    61.  
    62. class Bottleneck(nn.Module):
    63.     expansion = 2
    64.  
    65.     def __init__(self, inplanes, planes, stride=1, dilation=1):
    66.         super(Bottleneck, self).__init__()
    67.         expansion = Bottleneck.expansion
    68.         bottle_planes = planes // expansion
    69.         self.conv1 = nn.Conv2d(inplanes, bottle_planes,
    70.                                kernel_size=1, bias=False)
    71.         self.bn1 = nn.BatchNorm2d(bottle_planes, momentum=BN_MOMENTUM)
    72.         self.conv2 = nn.Conv2d(bottle_planes, bottle_planes, kernel_size=3,
    73.                                stride=stride, padding=dilation,
    74.                                bias=False, dilation=dilation)
    75.         self.bn2 = nn.BatchNorm2d(bottle_planes, momentum=BN_MOMENTUM)
    76.         self.conv3 = nn.Conv2d(bottle_planes, planes,
    77.                                kernel_size=1, bias=False)
    78.         self.bn3 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
    79.         self.relu = nn.ReLU(inplace=True)
    80.         self.stride = stride
    81.  
    82.     def forward(self, x, residual=None):
    83.         if residual is None:
    84.             residual = x
    85.  
    86.         out = self.conv1(x)
    87.         out = self.bn1(out)
    88.         out = self.relu(out)
    89.  
    90.         out = self.conv2(out)
    91.         out = self.bn2(out)
    92.         out = self.relu(out)
    93.  
    94.         out = self.conv3(out)
    95.         out = self.bn3(out)
    96.  
    97.         out += residual
    98.         out = self.relu(out)
    99.  
    100.         return out
    101.  
    102.  
    103. class BottleneckX(nn.Module):
    104.     expansion = 2
    105.     cardinality = 32
    106.  
    107.     def __init__(self, inplanes, planes, stride=1, dilation=1):
    108.         super(BottleneckX, self).__init__()
    109.         cardinality = BottleneckX.cardinality
    110.         # dim = int(math.floor(planes * (BottleneckV5.expansion / 64.0)))
    111.         # bottle_planes = dim * cardinality
    112.         bottle_planes = planes * cardinality // 32
    113.         self.conv1 = nn.Conv2d(inplanes, bottle_planes,
    114.                                kernel_size=1, bias=False)
    115.         self.bn1 = nn.BatchNorm2d(bottle_planes, momentum=BN_MOMENTUM)
    116.         self.conv2 = nn.Conv2d(bottle_planes, bottle_planes, kernel_size=3,
    117.                                stride=stride, padding=dilation, bias=False,
    118.                                dilation=dilation, groups=cardinality)
    119.         self.bn2 = nn.BatchNorm2d(bottle_planes, momentum=BN_MOMENTUM)
    120.         self.conv3 = nn.Conv2d(bottle_planes, planes,
    121.                                kernel_size=1, bias=False)
    122.         self.bn3 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
    123.         self.relu = nn.ReLU(inplace=True)
    124.         self.stride = stride
    125.  
    126.     def forward(self, x, residual=None):
    127.         if residual is None:
    128.             residual = x
    129.  
    130.         out = self.conv1(x)
    131.         out = self.bn1(out)
    132.         out = self.relu(out)
    133.  
    134.         out = self.conv2(out)
    135.         out = self.bn2(out)
    136.         out = self.relu(out)
    137.  
    138.         out = self.conv3(out)
    139.         out = self.bn3(out)
    140.  
    141.         out += residual
    142.         out = self.relu(out)
    143.  
    144.         return out
    145.  
    146.  
    147. class Root(nn.Module):
    148.     def __init__(self, in_channels, out_channels, kernel_size, residual):
    149.         super(Root, self).__init__()
    150.         self.conv = nn.Conv2d(
    151.             in_channels, out_channels, 1,
    152.             stride=1, bias=False, padding=(kernel_size - 1) // 2)
    153.         self.bn = nn.BatchNorm2d(out_channels, momentum=BN_MOMENTUM)
    154.         self.relu = nn.ReLU(inplace=True)
    155.         self.residual = residual
    156.  
    157.     def forward(self, *x):
    158.         children = x
    159.         x = self.conv(torch.cat(x, 1))
    160.         x = self.bn(x)
    161.         if self.residual:
    162.             x += children[0]
    163.         x = self.relu(x)
    164.  
    165.         return x
    166.  
    167.  
    168. class Tree(nn.Module):
    169.     def __init__(self, levels, block, in_channels, out_channels, stride=1,
    170.                  level_root=False, root_dim=0, root_kernel_size=1,
    171.                  dilation=1, root_residual=False):
    172.         super(Tree, self).__init__()
    173.         if root_dim == 0:
    174.             root_dim = 2 * out_channels
    175.         if level_root:
    176.             root_dim += in_channels
    177.         if levels == 1:
    178.             self.tree1 = block(in_channels, out_channels, stride,
    179.                                dilation=dilation)
    180.             self.tree2 = block(out_channels, out_channels, 1,
    181.                                dilation=dilation)
    182.         else:
    183.             self.tree1 = Tree(levels - 1, block, in_channels, out_channels,
    184.                               stride, root_dim=0,
    185.                               root_kernel_size=root_kernel_size,
    186.                               dilation=dilation, root_residual=root_residual)
    187.             self.tree2 = Tree(levels - 1, block, out_channels, out_channels,
    188.                               root_dim=root_dim + out_channels,
    189.                               root_kernel_size=root_kernel_size,
    190.                               dilation=dilation, root_residual=root_residual)
    191.         if levels == 1:
    192.             self.root = Root(root_dim, out_channels, root_kernel_size,
    193.                              root_residual)
    194.         self.level_root = level_root
    195.         self.root_dim = root_dim
    196.         self.downsample = None
    197.         self.project = None
    198.         self.levels = levels
    199.         if stride > 1:
    200.             self.downsample = nn.MaxPool2d(stride, stride=stride)
    201.         if in_channels != out_channels:
    202.             self.project = nn.Sequential(
    203.                 nn.Conv2d(in_channels, out_channels,
    204.                           kernel_size=1, stride=1, bias=False),
    205.                 nn.BatchNorm2d(out_channels, momentum=BN_MOMENTUM)
    206.             )
    207.  
    208.     def forward(self, x, residual=None, children=None):
    209.         children = [] if children is None else children
    210.         bottom = self.downsample(x) if self.downsample else x
    211.         residual = self.project(bottom) if self.project else bottom
    212.         if self.level_root:
    213.             children.append(bottom)
    214.         x1 = self.tree1(x, residual)
    215.         if self.levels == 1:
    216.             x2 = self.tree2(x1)
    217.             x = self.root(x2, x1, *children)
    218.         else:
    219.             children.append(x1)
    220.             x = self.tree2(x1, children=children)
    221.         return x
    222.  
    223.  
    224. class DLA(nn.Module):
    225.     def __init__(self, levels, channels, num_classes=1000,
    226.                  block=BasicBlock, residual_root=False, linear_root=False):
    227.         super(DLA, self).__init__()
    228.         self.channels = channels
    229.         self.num_classes = num_classes
    230.         self.base_layer = nn.Sequential(
    231.             nn.Conv2d(3, channels[0], kernel_size=7, stride=1,
    232.                       padding=3, bias=False),
    233.             nn.BatchNorm2d(channels[0], momentum=BN_MOMENTUM),
    234.             nn.ReLU(inplace=True))
    235.         self.level0 = self._make_conv_level(
    236.             channels[0], channels[0], levels[0])
    237.         self.level1 = self._make_conv_level(
    238.             channels[0], channels[1], levels[1], stride=2)
    239.         self.level2 = Tree(levels[2], block, channels[1], channels[2], 2,
    240.                            level_root=False,
    241.                            root_residual=residual_root)
    242.         self.level3 = Tree(levels[3], block, channels[2], channels[3], 2,
    243.                            level_root=True, root_residual=residual_root)
    244.         self.level4 = Tree(levels[4], block, channels[3], channels[4], 2,
    245.                            level_root=True, root_residual=residual_root)
    246.         self.level5 = Tree(levels[5], block, channels[4], channels[5], 2,
    247.                            level_root=True, root_residual=residual_root)
    248.  
    249.         # for m in self.modules():
    250.         #     if isinstance(m, nn.Conv2d):
    251.         #         n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
    252.         #         m.weight.data.normal_(0, math.sqrt(2. / n))
    253.         #     elif isinstance(m, nn.BatchNorm2d):
    254.         #         m.weight.data.fill_(1)
    255.         #         m.bias.data.zero_()
    256.  
    257.     def _make_level(self, block, inplanes, planes, blocks, stride=1):
    258.         downsample = None
    259.         if stride != 1 or inplanes != planes:
    260.             downsample = nn.Sequential(
    261.                 nn.MaxPool2d(stride, stride=stride),
    262.                 nn.Conv2d(inplanes, planes,
    263.                           kernel_size=1, stride=1, bias=False),
    264.                 nn.BatchNorm2d(planes, momentum=BN_MOMENTUM),
    265.             )
    266.  
    267.         layers = []
    268.         layers.append(block(inplanes, planes, stride, downsample=downsample))
    269.         for i in range(1, blocks):
    270.             layers.append(block(inplanes, planes))
    271.  
    272.         return nn.Sequential(*layers)
    273.  
    274.     def _make_conv_level(self, inplanes, planes, convs, stride=1, dilation=1):
    275.         modules = []
    276.         for i in range(convs):
    277.             modules.extend([
    278.                 nn.Conv2d(inplanes, planes, kernel_size=3,
    279.                           stride=stride if i == 0 else 1,
    280.                           padding=dilation, bias=False, dilation=dilation),
    281.                 nn.BatchNorm2d(planes, momentum=BN_MOMENTUM),
    282.                 nn.ReLU(inplace=True)])
    283.             inplanes = planes
    284.         return nn.Sequential(*modules)
    285.  
    286.     def forward(self, x):
    287.         y = []
    288.         x = self.base_layer(x)
    289.         for i in range(6):
    290.             x = getattr(self, 'level{}'.format(i))(x)
    291.             y.append(x)
    292.         return y
    293.  
    294.     def load_pretrained_model(self, data='imagenet', name='dla34', hash='ba72cf86'):
    295.         # fc = self.fc
    296.         if name.endswith('.pth'):
    297.             model_weights = torch.load(data + name)
    298.         else:
    299.             model_url = get_model_url(data, name, hash)
    300.             model_weights = model_zoo.load_url(model_url)
    301.         num_classes = len(model_weights[list(model_weights.keys())[-1]])
    302.         self.fc = nn.Conv2d(
    303.             self.channels[-1], num_classes,
    304.             kernel_size=1, stride=1, padding=0, bias=True)
    305.         self.load_state_dict(model_weights)
    306.         # self.fc = fc
    307.  
    308.  
    309. def dla34(pretrained=True, **kwargs):  # DLA-34
    310.     model = DLA([1, 1, 1, 2, 2, 1],
    311.                 [16, 32, 64, 128, 256, 512],
    312.                 block=BasicBlock, **kwargs)
    313.     #if pretrained:
    314.     #    model.load_pretrained_model(data='imagenet', name='dla34', hash='ba72cf86')
    315.     return model
    316.  
    317. class Identity(nn.Module):
    318.  
    319.     def __init__(self):
    320.         super(Identity, self).__init__()
    321.  
    322.     def forward(self, x):
    323.         return x
    324.  
    325.  
    326. def fill_fc_weights(layers):
    327.     for m in layers.modules():
    328.         if isinstance(m, nn.Conv2d):
    329.             if m.bias is not None:
    330.                 nn.init.constant_(m.bias, 0)
    331.  
    332.  
    333. def fill_up_weights(up):
    334.     w = up.weight.data
    335.     f = math.ceil(w.size(2) / 2)
    336.     c = (2 * f - 1 - f % 2) / (2. * f)
    337.     for i in range(w.size(2)):
    338.         for j in range(w.size(3)):
    339.             w[0, 0, i, j] = \
    340.                 (1 - math.fabs(i / f - c)) * (1 - math.fabs(j / f - c))
    341.     for c in range(1, w.size(0)):
    342.         w[c, 0, :, :] = w[0, 0, :, :]
    343.  
    344.  
    345. class DeformConv(nn.Module):
    346.     def __init__(self, chi, cho):
    347.         super(DeformConv, self).__init__()
    348.         self.actf = nn.Sequential(
    349.             nn.BatchNorm2d(cho, momentum=BN_MOMENTUM),
    350.             nn.ReLU(inplace=True)
    351.         )
    352.         self.conv = DCN(chi, cho, kernel_size=(3,3), stride=1, padding=1, dilation=1, deformable_groups=1)
    353.  
    354.     def forward(self, x):
    355.         x = self.conv(x)
    356.         x = self.actf(x)
    357.         return x
    358.  
    359.  
    360. class IDAUp(nn.Module):
    361.  
    362.     def __init__(self, o, channels, up_f):
    363.         super(IDAUp, self).__init__()
    364.         for i in range(1, len(channels)):
    365.             c = channels[i]
    366.             f = int(up_f[i])  
    367.             proj = DeformConv(c, o)
    368.             node = DeformConv(o, o)
    369.      
    370.             up = nn.ConvTranspose2d(o, o, f * 2, stride=f, 
    371.                                     padding=f // 2, output_padding=0,
    372.                                     groups=o, bias=False)
    373.             fill_up_weights(up)
    374.  
    375.             setattr(self, 'proj_' + str(i), proj)
    376.             setattr(self, 'up_' + str(i), up)
    377.             setattr(self, 'node_' + str(i), node)
    378.                  
    379.         
    380.     def forward(self, layers, startp, endp):
    381.         for i in range(startp + 1, endp):
    382.             upsample = getattr(self, 'up_' + str(i - startp))
    383.             project = getattr(self, 'proj_' + str(i - startp))
    384.             layers[i] = upsample(project(layers[i]))
    385.             node = getattr(self, 'node_' + str(i - startp))
    386.             layers[i] = node(layers[i] + layers[i - 1])
    387.  
    388.  
    389.  
    390. class DLAUp(nn.Module):
    391.     def __init__(self, startp, channels, scales, in_channels=None):
    392.         super(DLAUp, self).__init__()
    393.         self.startp = startp
    394.         if in_channels is None:
    395.             in_channels = channels
    396.         self.channels = channels
    397.         channels = list(channels)
    398.         scales = np.array(scales, dtype=int)
    399.         for i in range(len(channels) - 1):
    400.             j = -i - 2
    401.             setattr(self, 'ida_{}'.format(i),
    402.                     IDAUp(channels[j], in_channels[j:],
    403.                           scales[j:] // scales[j]))
    404.             scales[j + 1:] = scales[j]
    405.             in_channels[j + 1:] = [channels[j] for _ in channels[j + 1:]]
    406.  
    407.     def forward(self, layers):
    408.         out = [layers[-1]] # start with 32
    409.         for i in range(len(layers) - self.startp - 1):
    410.             ida = getattr(self, 'ida_{}'.format(i))
    411.             ida(layers, len(layers) -i - 2, len(layers))
    412.             out.insert(0, layers[-1])
    413.         return out
    414.  
    415.  
    416. class Interpolate(nn.Module):
    417.     def __init__(self, scale, mode):
    418.         super(Interpolate, self).__init__()
    419.         self.scale = scale
    420.         self.mode = mode
    421.         
    422.     def forward(self, x):
    423.         x = F.interpolate(x, scale_factor=self.scale, mode=self.mode, align_corners=False)
    424.         return x
    425.  
    426.  
    427. class DLASeg(nn.Module):
    428.     def __init__(self, base_name, heads, pretrained, down_ratio, final_kernel,
    429.                  last_level, head_conv, out_channel=0):
    430.         super(DLASeg, self).__init__()
    431.         assert down_ratio in [2, 4, 8, 16]
    432.         self.first_level = int(np.log2(down_ratio))
    433.         self.last_level = last_level
    434.         self.base = globals()[base_name](pretrained=pretrained)
    435.         channels = self.base.channels
    436.         scales = [2 ** i for i in range(len(channels[self.first_level:]))]
    437.         self.dla_up = DLAUp(self.first_level, channels[self.first_level:], scales)
    438.  
    439.         if out_channel == 0:
    440.             out_channel = channels[self.first_level]
    441.  
    442.         self.ida_up = IDAUp(out_channel, channels[self.first_level:self.last_level], 
    443.                             [2 ** i for i in range(self.last_level - self.first_level)])
    444.         
    445.         self.heads = heads
    446.         for head in self.heads:
    447.             classes = self.heads[head]
    448.             if head_conv > 0:
    449.               fc = nn.Sequential(
    450.                   nn.Conv2d(channels[self.first_level], head_conv,
    451.                     kernel_size=3, padding=1, bias=True),
    452.                   nn.ReLU(inplace=True),
    453.                   nn.Conv2d(head_conv, classes, 
    454.                     kernel_size=final_kernel, stride=1, 
    455.                     padding=final_kernel // 2, bias=True))
    456.               if 'hm' in head:
    457.                 fc[-1].bias.data.fill_(-2.19)
    458.               else:
    459.                 fill_fc_weights(fc)
    460.             else:
    461.               fc = nn.Conv2d(channels[self.first_level], classes, 
    462.                   kernel_size=final_kernel, stride=1, 
    463.                   padding=final_kernel // 2, bias=True)
    464.               if 'hm' in head:
    465.                 fc.bias.data.fill_(-2.19)
    466.               else:
    467.                 fill_fc_weights(fc)
    468.             self.__setattr__(head, fc)
    469.  
    470.     def forward(self, x):
    471.         x = self.base(x)
    472.         x = self.dla_up(x)
    473.  
    474.         y = []
    475.         for i in range(self.last_level - self.first_level):
    476.             y.append(x[i].clone())
    477.         self.ida_up(y, 0, len(y))
    478.  
    479.         z = {}
    480.         for head in self.heads:
    481.             z[head] = self.__getattr__(head)(y[-1])
    482.         # return [z]  xiehao
    483.         return z
    484.     
    485.  
    486. def get_pose_net(num_layers, heads, head_conv=256, down_ratio=4):
    487.   model = DLASeg('dla{}'.format(num_layers), heads,
    488.                  pretrained=True,
    489.                  down_ratio=down_ratio,
    490.                  final_kernel=1,
    491.                  last_level=5,
    492.                  head_conv=head_conv)
    493.   return model
    494.

    测试剪枝主流程的centerNet_prune文件

    1. from nni.compression.pytorch.pruning import L1NormPruner
    2. from nni.compression.pytorch.speedup import ModelSpeedup
    3. import pose_dla_dcn
    4. import torch
    5.  
    6. num_layers = 34
    7. heads = {'hm': 2, 'wh': 2, 'reg': 2}
    8. config_list = [{
    9.     'sparsity': 0.5,
    10.     'op_types': ['Conv2d']
    11. }]
    12.  
    13. model = pose_dla_dcn.get_pose_net(num_layers, heads)
    14. torch.save(model.state_dict(), "baseline.pth")
    15. # model.cuda()
    16. # model(torch.rand(1, 3, 512, 512).cuda())  # 512严格要求
    17.  
    18. print("--------------raw model--------------")
    19. print(model)
    20.  
    21. pruner = L1NormPruner(model, config_list, mode='dependency_aware', dummy_input=torch.rand(1, 3, 512, 512))
    22. print("--------------pruned model--------------")
    23. print(model)
    24.  
    25. # compress the model and generate the masks
    26. _, masks = pruner.compress()
    27. # show the masks sparsity
    28. for name, mask in masks.items():
    29.     print(name, ' sparsity: ', '{:.2f}'.format(mask['weight'].sum() / mask['weight'].numel()))
    30.  
    31. # need to unwarp the model, if the model is wrawpped before speedup
    32. pruner._unwrap_model()
    33.  
    34. # speedup the model
    35. model.eval()
    36. ModelSpeedup(model, dummy_input=torch.rand(1, 3, 512, 512), masks_file=masks).speedup_model()
    37.  
    38. print("--------------after sppedup--------------")
    39. print(model)
    40. torch.save(model, "model_pruner.pth")

    最后一个文件有一个注意点，就是dummy_input输入一定是[batch, 3, 512, 512]的形式，否则卷积池化后残差连接会有问题，会出现[batch, 128, 3, 3] 和[batch, 128, 2, 2]相加的情况。

    执行后报错，出错信息为：

    self.output = self.module(*dummy_input)
      File "D:\programs\python37\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl
        result = self.forward(*input, **kwargs)
    TypeError: forward() missing 1 required positional argument: 'input'

    1. [2022-11-18 19:15:08] start to speedup the model
    2. D:\programs\python37\lib\site-packages\nni\compression\pytorch\utils\mask_conflict.py:124: UserWarning: This overload of nonzero is deprecated:
    3. 	nonzero()
    4. Consider using one of the following signatures instead:
    5. 	nonzero(*, bool as_tuple) (Triggered internally at  ..\torch\csrc\utils\python_arg_parser.cpp:882.)
    6.   all_ones = (w_mask.flatten(1).sum(-1) == count).nonzero().squeeze(1).tolist()
    7. both dim0 and dim1 masks found.
    8. [2022-11-18 19:18:01] infer module masks...
    9. [2022-11-18 19:18:01] Update mask for base.base_layer.0
    10. [2022-11-18 19:18:03] Update mask for dla_up.ida_0.proj_1.actf.0
    11. Traceback (most recent call last):
    12.   File "D:/workspace/newczalgo/CenterNet_3DCar_Slimming/src/centernet_prune/centernet_prune.py", line 34, in 
    13.     ModelSpeedup(model, torch.rand(1, 3, 512, 512), masks).speedup_model()
    14.   File "D:\programs\python37\lib\site-packages\nni\compression\pytorch\speedup\compressor.py", line 536, in speedup_model
    15.     self.infer_modules_masks()
    16.   File "D:\programs\python37\lib\site-packages\nni\compression\pytorch\speedup\compressor.py", line 371, in infer_modules_masks
    17.     self.update_direct_sparsity(curnode)
    18.   File "D:\programs\python37\lib\site-packages\nni\compression\pytorch\speedup\compressor.py", line 240, in update_direct_sparsity
    19.     state_dict=copy.deepcopy(module.state_dict()), batch_dim=self.batch_dim)
    20.   File "D:\programs\python37\lib\site-packages\nni\compression\pytorch\speedup\infer_mask.py", line 80, in __init__
    21.     self.output = self.module(*dummy_input)
    22.   File "D:\programs\python37\lib\site-packages\torch\nn\modules\module.py", line 727, in _call_impl
    23.     result = self.forward(*input, **kwargs)
    24. TypeError: forward() missing 1 required positional argument: 'input'
    25.  
    26. Process finished with exit code 1

    nni的github上有对应的issue：

    Meet a dropout missing argument problem when speed up the model and update mask for my own model · Issue #4297 · microsoft/nni · GitHub

     主要是DCN的输出为empy，导致dummy_input的list为空的。

    尝试方法一：

    网上说把有问题的层信息加入config_list的exclude中，比如:

    1. config_list = [{
    2.     'sparsity': 0.5,
    3.     'op_types': ['Conv2d']
    4. },
    5. {'exclude': True, 'op_names': ['conv_offset_mask']}
    6. ]

     op_names换成：

    dla_up.ida_0.proj_1.actf.0 或 dla_up.ida_0.proj_1.conv.conv_offset_mask或dla_up等都失败，报错信息还是一样的。

    尝试方法二：

    在compress中当dumpy_input为空时直接return，如下所示：

    此时又会出现新的bug：

    1. [2022-11-18 21:43:10] Update mask for dla_up.ida_2.aten::add.233
    2. Traceback (most recent call last):
    3.   File "D:/workspace/newczalgo/CenterNet_3DCar_Slimming/src/centernet_prune/centernet_prune.py", line 38, in 
    4.     ModelSpeedup(model, dummy_input=torch.rand(1, 3, 512, 512), masks_file=masks).speedup_model()
    5.   File "D:\programs\python37\lib\site-packages\nni\compression\pytorch\speedup\compressor.py", line 539, in speedup_model
    6.     self.infer_modules_masks()
    7.   File "D:\programs\python37\lib\site-packages\nni\compression\pytorch\speedup\compressor.py", line 374, in infer_modules_masks
    8.     self.update_direct_sparsity(curnode)
    9.   File "D:\programs\python37\lib\site-packages\nni\compression\pytorch\speedup\compressor.py", line 235, in update_direct_sparsity
    10.     func, dummy_input, in_masks, in_constants=in_constants, batch_dim=self.batch_dim)
    11.   File "D:\programs\python37\lib\site-packages\nni\compression\pytorch\speedup\infer_mask.py", line 80, in __init__
    12.     self.output = self.module(*dummy_input)
    13. TypeError: add() received an invalid combination of arguments - got (Tensor), but expected (Tensor input, Tensor other, *, Number alpha, Tensor out)

    结论：DLASeg使用NNI进行剪枝挑战失败！！

  • 相关阅读:
    前端ps基本操作
    滚动播报、el-scrollbar
    C语言解决约瑟夫环问题
    R语言ggplot2和gganimate包可视化动态动画气泡图（Animated Bubble chart）：使用gganimate包创建可视化gif动图
    flink源码分析 - 获取调用位置信息
    vscode中注释多行bash脚本
    远程控制 Linux 系统的软件下载
    gdb-dashboard的简单使用
    到底是选择极米还是选择当贝？画质看当贝，音质看极米 在意画质和使用体验的选当贝
    Vscode连接远程服务器（一套配置成功）
  • 原文地址：https://blog.csdn.net/benben044/article/details/127931805