fix export (#9197)

Author: Birdylx

ppdet/engine/trainer.py

@@ -1266,6 +1266,10 @@ def export(self, output_dir='output_inference', for_fd=False):
         self.model.eval()
         model = copy.deepcopy(self.model)
 
+        # add `export_mode` attr for all layers
+        for layer in self.model.sublayers(include_self=True):
+            layer.export_mode = True
+
         model_name = os.path.splitext(os.path.split(self.cfg.filename)[-1])[0]
         if for_fd:
             save_dir = output_dir

ppdet/modeling/heads/ppyoloe_ins_head.py

@@ -686,6 +686,10 @@ def post_process(self,
                      scale_factor,
                      infer_shape=[640, 640],
                      rescale=True):
+
+        if getattr(self, "export_mode", False):
+            return self.export_post_process(head_outs, im_shape, scale_factor, infer_shape, rescale)
+
         pred_scores, pred_dist, pred_mask_coeffs, mask_feat, anchor_points, stride_tensor = head_outs
 
         pred_bboxes = batch_distance2bbox(anchor_points, pred_dist)
@@ -740,4 +744,55 @@ def post_process(self,
             bbox_pred = paddle.zeros([bbox_num, 6])
             mask_pred = paddle.zeros([bbox_num, int(ori_h), int(ori_w)])
 
+        return bbox_pred, bbox_num, mask_pred, keep_idxs
+    
+    def export_post_process(self,
+                            head_outs,
+                            im_shape,
+                            scale_factor,
+                            infer_shape=[640, 640],
+                            rescale=True):
+        pred_scores, pred_dist, pred_mask_coeffs, mask_feat, anchor_points, stride_tensor = head_outs
+
+        pred_bboxes = batch_distance2bbox(anchor_points, pred_dist)
+        pred_bboxes *= stride_tensor
+
+        if self.exclude_post_process:
+            return paddle.concat([
+                pred_bboxes,
+                pred_scores.transpose([0, 2, 1]),
+                pred_mask_coeffs.transpose([0, 2, 1])
+            ],
+                                 axis=-1), mask_feat, None
+            # [1, 8400, 4+80+32], [1, 32, 160, 160]
+
+        bbox_pred, bbox_num, keep_idxs = self.nms(pred_bboxes, pred_scores)
+
+        pred_mask_coeffs = pred_mask_coeffs.transpose([0, 2, 1])
+        mask_coeffs = paddle.gather(
+            pred_mask_coeffs.reshape([-1, self.num_masks]), keep_idxs)
+
+        mask_logits = process_mask_upsample(mask_feat[0], mask_coeffs,
+                                            bbox_pred[:, 2:6], infer_shape)
+        if rescale:
+            ori_h, ori_w = im_shape[0] / scale_factor[0]
+            mask_logits = F.interpolate(
+                mask_logits.unsqueeze(0),
+                size=[
+                    int(paddle.round(mask_logits.shape[-2] /
+                            scale_factor[0][0])),
+                    int(paddle.round(mask_logits.shape[-1] /
+                            scale_factor[0][1]))
+                ],
+                mode='bilinear',
+                align_corners=False)
+            mask_logits = mask_logits[..., :int(ori_h), :int(ori_w)]
+
+        masks = mask_logits.squeeze(0)
+        mask_pred = masks > self.mask_thr_binary
+
+        # scale bbox to origin
+        scale_factor = scale_factor.flip(-1).tile([1, 2])
+        bbox_pred[:, 2:6] /= scale_factor
+
         return bbox_pred, bbox_num, mask_pred, keep_idxs