From 750ea214c6c252527e7f6f24a71f34d569690765 Mon Sep 17 00:00:00 2001 From: Guanghua Yu <742925032@qq.com> Date: Wed, 16 Mar 2022 23:11:57 +0800 Subject: [PATCH] Support nms options when export model (#5373) * support nms options when export model * fix config --- configs/picodet/_base_/picodet_esnet.yml | 1 - configs/picodet/_base_/picodet_esnetv2.yml | 1 - configs/runtime.yml | 6 ++++ ppdet/engine/export_utils.py | 5 +-- ppdet/engine/trainer.py | 21 +++++++++---- ppdet/modeling/architectures/picodet.py | 11 +++++-- ppdet/modeling/heads/pico_head.py | 36 ++++++++++++++++++++++ 7 files changed, 68 insertions(+), 13 deletions(-) diff --git a/configs/picodet/_base_/picodet_esnet.yml b/configs/picodet/_base_/picodet_esnet.yml index 24b213f2a..aa099fca1 100644 --- a/configs/picodet/_base_/picodet_esnet.yml +++ b/configs/picodet/_base_/picodet_esnet.yml @@ -1,6 +1,5 @@ architecture: PicoDet pretrain_weights: https://paddledet.bj.bcebos.com/models/pretrained/ESNet_x1_0_pretrained.pdparams -export_post_process: False # Whether post-processing is included in the network when export model. PicoDet: backbone: ESNet diff --git a/configs/picodet/_base_/picodet_esnetv2.yml b/configs/picodet/_base_/picodet_esnetv2.yml index 8e16d5298..eab2d4115 100644 --- a/configs/picodet/_base_/picodet_esnetv2.yml +++ b/configs/picodet/_base_/picodet_esnetv2.yml @@ -1,6 +1,5 @@ architecture: PicoDet pretrain_weights: https://paddledet.bj.bcebos.com/models/pretrained/ESNet_x1_0_pretrained.pdparams -export_post_process: False # Whether post-processing is included in the network when export model. PicoDet: backbone: ESNet diff --git a/configs/runtime.yml b/configs/runtime.yml index 273b8b1af..c67c6c94f 100644 --- a/configs/runtime.yml +++ b/configs/runtime.yml @@ -4,3 +4,9 @@ log_iter: 20 save_dir: output snapshot_epoch: 1 print_flops: false + +# Exporting the model +export: + post_process: True # Whether post-processing is included in the network when export model. + nms: True # Whether NMS is included in the network when export model. + benchmark: False # It is used to testing model performance, if set `True`, post-process and NMS will not be exported. diff --git a/ppdet/engine/export_utils.py b/ppdet/engine/export_utils.py index 9742d9665..b73b77e8b 100644 --- a/ppdet/engine/export_utils.py +++ b/ppdet/engine/export_utils.py @@ -41,7 +41,7 @@ TRT_MIN_SUBGRAPH = { 'HigherHRNet': 3, 'HRNet': 3, 'DeepSORT': 3, - 'ByteTrack':10, + 'ByteTrack': 10, 'JDE': 10, 'FairMOT': 5, 'GFL': 16, @@ -166,7 +166,8 @@ def _dump_infer_config(config, path, image_shape, model): reader_cfg, dataset_cfg, config['metric'], label_arch, image_shape[1:]) if infer_arch == 'PicoDet': - if config.get('export_post_process', False): + if hasattr(config, 'export') and config['export'].get('post_process', + False): infer_cfg['arch'] = 'GFL' head_name = 'PicoHeadV2' if config['PicoHeadV2'] else 'PicoHead' infer_cfg['NMS'] = config[head_name]['nms'] diff --git a/ppdet/engine/trainer.py b/ppdet/engine/trainer.py index 2573ada01..f186afe8e 100644 --- a/ppdet/engine/trainer.py +++ b/ppdet/engine/trainer.py @@ -651,13 +651,21 @@ class Trainer(object): if hasattr(layer, 'convert_to_deploy'): layer.convert_to_deploy() - export_post_process = self.cfg.get('export_post_process', False) - if hasattr(self.model, 'export_post_process'): - self.model.export_post_process = export_post_process - image_shape = [None] + image_shape[1:] + export_post_process = self.cfg['export'].get( + 'post_process', False) if hasattr(self.cfg, 'export') else True + export_nms = self.cfg['export'].get('nms', False) if hasattr( + self.cfg, 'export') else True + export_benchmark = self.cfg['export'].get( + 'benchmark', False) if hasattr(self.cfg, 'export') else False if hasattr(self.model, 'fuse_norm'): self.model.fuse_norm = self.cfg['TestReader'].get('fuse_normalize', False) + if hasattr(self.model, 'export_post_process'): + self.model.export_post_process = export_post_process if not export_benchmark else False + if hasattr(self.model, 'export_nms'): + self.model.export_nms = export_nms if not export_benchmark else False + if export_post_process and not export_benchmark: + image_shape = [None] + image_shape[1:] # Save infer cfg _dump_infer_config(self.cfg, @@ -789,5 +797,6 @@ class Trainer(object): images.sort() assert len(images) > 0, "no image found in {}".format(infer_dir) all_images.extend(images) - logger.info("Found {} inference images in total.".format(len(images))) - return all_images \ No newline at end of file + logger.info("Found {} inference images in total.".format( + len(images))) + return all_images diff --git a/ppdet/modeling/architectures/picodet.py b/ppdet/modeling/architectures/picodet.py index 7331fb712..760b8347b 100644 --- a/ppdet/modeling/architectures/picodet.py +++ b/ppdet/modeling/architectures/picodet.py @@ -42,6 +42,7 @@ class PicoDet(BaseArch): self.neck = neck self.head = head self.export_post_process = True + self.export_nms = True @classmethod def from_config(cls, cfg, *args, **kwargs): @@ -68,8 +69,8 @@ class PicoDet(BaseArch): else: im_shape = self.inputs['im_shape'] scale_factor = self.inputs['scale_factor'] - bboxes, bbox_num = self.head.post_process(head_outs, im_shape, - scale_factor) + bboxes, bbox_num = self.head.post_process( + head_outs, im_shape, scale_factor, export_nms=self.export_nms) return bboxes, bbox_num def get_loss(self, ): @@ -85,7 +86,11 @@ class PicoDet(BaseArch): def get_pred(self): if not self.export_post_process: return {'picodet': self._forward()[0]} - else: + elif self.export_nms: bbox_pred, bbox_num = self._forward() output = {'bbox': bbox_pred, 'bbox_num': bbox_num} return output + else: + bboxes, mlvl_scores = self._forward() + output = {'bbox': bboxes, 'scores': mlvl_scores} + return output diff --git a/ppdet/modeling/heads/pico_head.py b/ppdet/modeling/heads/pico_head.py index 8a5e04412..98c8c8ef9 100644 --- a/ppdet/modeling/heads/pico_head.py +++ b/ppdet/modeling/heads/pico_head.py @@ -335,6 +335,24 @@ class PicoHead(OTAVFLHead): return (cls_logits_list, bboxes_reg_list) + def post_process(self, + gfl_head_outs, + im_shape, + scale_factor, + export_nms=True): + cls_scores, bboxes_reg = gfl_head_outs + bboxes = paddle.concat(bboxes_reg, axis=1) + mlvl_scores = paddle.concat(cls_scores, axis=1) + mlvl_scores = mlvl_scores.transpose([0, 2, 1]) + if not export_nms: + return bboxes, mlvl_scores + else: + # rescale: [h_scale, w_scale] -> [w_scale, h_scale, w_scale, h_scale] + im_scale = scale_factor.flip([1]).tile([1, 2]).unsqueeze(1) + bboxes /= im_scale + bbox_pred, bbox_num, _ = self.nms(bboxes, mlvl_scores) + return bbox_pred, bbox_num + @register class PicoHeadV2(GFLHead): @@ -625,3 +643,21 @@ class PicoHeadV2(GFLHead): loss_vfl=loss_vfl, loss_bbox=loss_bbox, loss_dfl=loss_dfl) return loss_states + + def post_process(self, + gfl_head_outs, + im_shape, + scale_factor, + export_nms=True): + cls_scores, bboxes_reg = gfl_head_outs + bboxes = paddle.concat(bboxes_reg, axis=1) + mlvl_scores = paddle.concat(cls_scores, axis=1) + mlvl_scores = mlvl_scores.transpose([0, 2, 1]) + if not export_nms: + return bboxes, mlvl_scores + else: + # rescale: [h_scale, w_scale] -> [w_scale, h_scale, w_scale, h_scale] + im_scale = scale_factor.flip([1]).tile([1, 2]).unsqueeze(1) + bboxes /= im_scale + bbox_pred, bbox_num, _ = self.nms(bboxes, mlvl_scores) + return bbox_pred, bbox_num -- GitLab