Merge pull request #235 from Birdylx/fix_export

Fix ppyoloe_seg export

Author: lyuwenyu

ppdet/engine/trainer.py

@@ -1183,6 +1183,10 @@ def export(self, output_dir='output_inference'):
         if hasattr(self.cfg, 'export') and 'fuse_conv_bn' in self.cfg[
                 'export'] and self.cfg['export']['fuse_conv_bn']:
             self.model = fuse_conv_bn(self.model)
+        
+        # enable export_mode
+        for layer in self.model.sublayers(include_self=True):
+                layer.in_export_mode = True
 
         model_name = os.path.splitext(os.path.split(self.cfg.filename)[-1])[0]
         save_dir = os.path.join(output_dir, model_name)

ppdet/modeling/heads/ppyoloe_ins_head.py

@@ -613,8 +613,11 @@ def post_process(self,
                         custom_ceil(mask_logits.shape[-1] / scale_factor[0][1])
                     ],
                     mode='bilinear',
-                    align_corners=False)[..., :round(ori_h.item()), :round(ori_w.item())] # due to npu numeric error, we need to take round of img size.
-                    # align_corners=False)[..., :int(ori_h), :int(ori_w)] # TODO: only for export
+                    align_corners=False)
+                if self.in_export_mode:
+                    mask_logits = mask_logits[..., :int(ori_h), :int(ori_w)]
+                else:
+                    mask_logits = mask_logits[..., :round(ori_h.item()), :round(ori_w.item())] # due to npu numeric error, we need to take round of img size.
             masks = mask_logits.squeeze(0)
             mask_pred = masks > self.mask_thr_binary