快捷模式在mmdetection中部署：

注意事项

mmdetection中对模型的forward 函数进行了重载，与一般的forward只需要输入Tensor不同，还需要img_meta，这与EasyPruner所需要的jit.trace里面的forward函数不兼容。本工具包参考了mmdetection自带的转onnx代码的一些写法，进行解决。

示例

以FOCS目标检测方法或者SCRFD人脸检测方法为例：

步骤1 常规训练

步骤2 准备微调代码，在微调前，加载预训练模型，并加入剪枝代码

之后，在建模型后，加载常规训练得到的模型权重，并进行剪枝，SCRFD中不支持默认加载预训练模型，需要手工加载一下，例如在mmdet/apis/train.py 的train_detector中修改如下：

  1.             data_loaders = [
  2.           build_dataloader(
  3.               ds,
  4.               cfg.data.samples_per_gpu,
  5.               cfg.data.workers_per_gpu,
  6.               # cfg.gpus will be ignored if distributed
  7.               len(cfg.gpu_ids),
  8.               dist=distributed,
  9.               seed=cfg.seed) for ds in dataset
  10.       ]
  11.   '''
  12.   以上为原始mmdetection代码
  13.   '''
  14.   # 对网络进行修改，使其改为单变量输入
  15.     from easypruner import  fastpruner
  16.     from easypruner.utils.rebuild import rebuild
  18.     from functools import partial
  19.     import copy
  20.     from mmdet.core.export import pytorch2onnx
  21.     imgpath = "data/retinaface/val/images/0--Parade/0_Parade_Parade_0_194.jpg"  # 需指定任意一个图片
  22.     normalize_cfg = {'mean': [127,127,127], 'std': [55,55,55]} #格式需要，随便指定即可
  23.     input_shape = [1,3,640,640] # 格式需要，按实际情况指定即可
  24.     input_config = { 
  25.         'input_shape': input_shape,
  26.         'input_path': imgpath,
  27.         'normalize_cfg': normalize_cfg
  28.     }   
  29.     one_img, one_meta = pytorch2onnx.preprocess_example_input(input_config)
  31.     model.forward = partial(
  32.         model.forward, img_metas=[[one_meta]], return_loss=False)
  35.     #加载模型并剪枝
  36.     state_dict = torch.load("34Gmodel.pth")
  37.     device = next(model.parameters()).device
  38.     model.load_state_dict(state_dict['state_dict'],strict=True)
  39.     model.cpu()
  41.     fastpruner.fastpruner(model, prune_factor = 0.4, method="Ratio", input_dim=[3,640,640])##Ratio 和uniform两种方式都可以试试，注意大小写
  42.     #fastpruner.fastpruner(model, prune_factor = 0.5, method="Uniform", input_dim=[3,640,640])##Ratio 和uniform两种方式都可以试试，注意大小写
  43.     #fastpruner.fastpruner(model, prune_factor = 0.01, method="Order", input_dim=[3,640,640])##Ratio 和uniform两种方式都可以试试，注意大小写
  44.     model.to(device)
  48.     #保存剪枝后的模型权重
  49.     state_dict = state_dict['state_dict']
  50.     for k,v in model.state_dict().items():
  51.         state_dict[k] = v
  52.     save_path = './model_pruned_34_0.4uniform.pt' #
  53.     torch.save(model.state_dict(),save_path) 
  54.     #torch.export()
  56.     #复原模型
  57.     model.forward = model.forward.func 
  58.     exit(0) #剪枝完可以直接训练finetune/或者exit退出，后续通过rebuild重构网络再finetune
  59.  '''
  60.  以下为原始mmdetection代码
  61.  '''
  62.      # put model on gpus
  63.       if distributed:
  64.          find_unused_parameters = cfg.get('find_unused_parameters', False)
  65.          # Sets the `find_unused_parameters` parameter in
  66.          # torch.nn.parallel.DistributedDataParallel
  67.          model = MMDistributedDataParallel(
  68.              model.cuda(),
  69.              device_ids=[torch.cuda.current_device()],
  70.              broadcast_buffers=False,
  71.              find_unused_parameters=find_unused_parameters)
  72.       else:
  73.          model = MMDataParallel(
  74.              model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)

使用uniform或者测试模型flops压缩率时，需要将模型定义与模型前向中(甚至模型转onnx）函数中的neck和head部分去掉。因为bbox每次产生的个数不一样所以flops也不一样，所以我们仅考虑backbone的flops大小。比如faster_rcnn，我们需要修改 mmdet/models/detectors/two_stage.py。以2.18.0版本mmdet为例：

  1. @DETECTORS.register_module()
  2. class TwoStageDetector(BaseDetector):
  3.     """Base class for two-stage detectors.
  5.     Two-stage detectors typically consisting of a region proposal network and a
  6.     task-specific regression head.
  7.     """
  9.     def __init__(self,
  10.                  backbone,
  11.                  neck=None,
  12.                  rpn_head=None,
  13.                  roi_head=None,
  14.                  train_cfg=None,
  15.                  test_cfg=None,
  16.                  pretrained=None,
  17.                  init_cfg=None):
  18.         super(TwoStageDetector, self).__init__(init_cfg)
  19.         if pretrained:
  20.             warnings.warn('DeprecationWarning: pretrained is deprecated, '
  21.                           'please use "init_cfg" instead')
  22.             backbone.pretrained = pretrained
  23.         self.backbone = build_backbone(backbone)
  24.         ####注释掉这里
  25.         '''
  26.         if neck is not None:
  27.             self.neck = build_neck(neck)
  29.         if rpn_head is not None:
  30.             rpn_train_cfg = train_cfg.rpn if train_cfg is not None else None
  31.             rpn_head_ = rpn_head.copy()
  32.             rpn_head_.update(train_cfg=rpn_train_cfg, test_cfg=test_cfg.rpn)
  33.             self.rpn_head = build_head(rpn_head_)
  35.         if roi_head is not None:
  36.             # update train and test cfg here for now
  37.             # TODO: refactor assigner & sampler
  38.             rcnn_train_cfg = train_cfg.rcnn if train_cfg is not None else None
  39.             roi_head.update(train_cfg=rcnn_train_cfg)
  40.             roi_head.update(test_cfg=test_cfg.rcnn)
  41.             roi_head.pretrained = pretrained
  42.             self.roi_head = build_head(roi_head)
  43.         '''
  44.         self.train_cfg = train_cfg
  45.         self.test_cfg = test_cfg
  47.     @property
  48.     def with_rpn(self):
  49.         """bool: whether the detector has RPN"""
  50.         return hasattr(self, 'rpn_head') and self.rpn_head is not None
  52. ...
  53.     def forward_dummy(self, img):
  54.         """Used for computing network flops.
  56.         See `mmdetection/tools/analysis_tools/get_flops.py`
  57.         """
  58.         outs = ()
  59.         # backbone
  60.         x = self.extract_feat(img)
  61.         return x #增加这行
  62.         #注释这里
  63.         '''
  64.         # rpn
  65.         if self.with_rpn:
  66.             rpn_outs = self.rpn_head(x)
  67.             outs = outs + (rpn_outs, )
  68.         proposals = torch.randn(1000, 4).to(img.device)
  69.         # roi_head
  70.         roi_outs = self.roi_head.forward_dummy(x, proposals)
  71.         outs = outs + (roi_outs, )
  72.         return outs
  73.         '''
  74.     def forward_train(self,
  75.                       img,
  76.                       img_metas,
  77.                       gt_bboxes,
  78.                       gt_labels,
  79.                       gt_bboxes_ignore=None,
  80.                       gt_masks=None,
  81.                       proposals=None,
  82.                       **kwargs):
  83.         """
  84.         Args:
  85.             img (Tensor): of shape (N, C, H, W) encoding input images.
  86.                 Typically these should be mean centered and std scaled.
  88.             img_metas (list[dict]): list of image info dict where each dict
  89.                 has: 'img_shape', 'scale_factor', 'flip', and may also contain
  90.                 'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
  91.                 For details on the values of these keys see
  92.                 `mmdet/datasets/pipelines/formatting.py:Collect`.
  94.             gt_bboxes (list[Tensor]): Ground truth bboxes for each image with
  95.                 shape (num_gts, 4) in [tl_x, tl_y, br_x, br_y] format.
  97.             gt_labels (list[Tensor]): class indices corresponding to each box
  99.             gt_bboxes_ignore (None | list[Tensor]): specify which bounding
  100.                 boxes can be ignored when computing the loss.
  102.             gt_masks (None | Tensor) : true segmentation masks for each box
  103.                 used if the architecture supports a segmentation task.
  105.             proposals : override rpn proposals with custom proposals. Use when
  106.                 `with_rpn` is False.
  108.         Returns:
  109.             dict[str, Tensor]: a dictionary of loss components
  110.         """
  111.         x = self.extract_feat(img)
  112.         return x #增加这行
  113.         #注释这里
  114.         '''
  115.         losses = dict()
  117.         # RPN forward and loss
  118.         if self.with_rpn:
  119.             proposal_cfg = self.train_cfg.get('rpn_proposal',
  120.                                               self.test_cfg.rpn)
  121.             rpn_losses, proposal_list = self.rpn_head.forward_train(
  122.                 x,
  123.                 img_metas,
  124.                 gt_bboxes,
  125.                 gt_labels=None,
  126.                 gt_bboxes_ignore=gt_bboxes_ignore,
  127.                 proposal_cfg=proposal_cfg,
  128.                 **kwargs)
  129.             losses.update(rpn_losses)
  130.         else:
  131.             proposal_list = proposals
  133.         roi_losses = self.roi_head.forward_train(x, img_metas, proposal_list,
  134.                                                  gt_bboxes, gt_labels,
  135.                                                  gt_bboxes_ignore, gt_masks,
  136.                                                  **kwargs)
  137.         losses.update(roi_losses)
  139.         return losses
  140.         '''
  141. ...
  142.    def simple_test(self, img, img_metas, proposals=None, rescale=False):
  143.         """Test without augmentation."""
  144.         #注释这里
  145.         #assert self.with_bbox, 'Bbox head must be implemented.'
  146.         x = self.extract_feat(img)
  147.         return x #增加这行
  148.         #注释这里
  149.         '''
  150.         if proposals is None:
  151.             proposal_list = self.rpn_head.simple_test_rpn(x, img_metas)
  152.         else:
  153.             proposal_list = proposals
  155.         return self.roi_head.simple_test(
  156.             x, proposal_list, img_metas, rescale=rescale)
  157.         '''
  158. ...
  159.     def onnx_export(self, img, img_metas):
  161.         img_shape = torch._shape_as_tensor(img)[2:]
  162.         img_metas[0]['img_shape_for_onnx'] = img_shape
  163.         x = self.extract_feat(img)
  164.         return x #增加这行
  165.         #注释这里
  166.         '''
  167.         proposals = self.rpn_head.onnx_export(x, img_metas)
  168.         if hasattr(self.roi_head, 'onnx_export'):
  169.             return self.roi_head.onnx_export(x, proposals, img_metas)
  170.         else:
  171.             raise NotImplementedError(
  172.                 f'{self.__class__.__name__} can not '
  173.                 f'be exported to ONNX. Please refer to the '
  174.                 f'list of supported models,'
  175.                 f'https://mmdetection.readthedocs.io/en/latest/tutorials/pytorch2onnx.html#list-of-supported-models-exportable-to-onnx'  # noqa E501
  176.             )
  177.         '''
  178. ...

步骤 3 基于常规训练的模型微调

步骤 4 单独准确率测试、部署转onnx时，需要将权重加载代码改为用rebuild函数加载，如以下代码：

依然在mmdet/apis/train.py 的train_detector中进行修改

            data_loaders = [
          build_dataloader(
              ds,
              cfg.data.samples_per_gpu,
              cfg.data.workers_per_gpu,
              # cfg.gpus will be ignored if distributed
              len(cfg.gpu_ids),
              dist=distributed,
              seed=cfg.seed) for ds in dataset
      ]
  '''
  以上为原始mmdetection代码
  '''        
    from easypruner.utils.rebuild import rebuild 
    state_dict = torch.load("model_pruned_34_0.4uniform_finetuned.pt")#剪枝后的模型权重文件
    #import pdb;pdb.set_trace() 
    model = rebuild(model , state_dict)
  '''
  以下为原始mmdetection代码
  '''
     # put model on gpus
      if distributed:
         find_unused_parameters = cfg.get('find_unused_parameters', False)
         # Sets the `find_unused_parameters` parameter in
         # torch.nn.parallel.DistributedDataParallel
         model = MMDistributedDataParallel(
             model.cuda(),
             device_ids=[torch.cuda.current_device()],
             broadcast_buffers=False,
             find_unused_parameters=find_unused_parameters)
      else:
         model = MMDataParallel(
             model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)