Add If-Else-Return of Yolov3 as Unittest (PaddlePaddle#25385)

As the title
MRXLT · Jul 7, 2020 · 30185ef · 30185ef
1 parent 80f1c50
commit 30185ef
Showing 1 changed file with 16 additions and 18 deletions.
diff --git a/python/paddle/fluid/tests/unittests/dygraph_to_static/yolov3.py b/python/paddle/fluid/tests/unittests/dygraph_to_static/yolov3.py
@@ -314,21 +314,19 @@ def forward(self,
                         scores, perm=[0, 2, 1]))
             self.downsample //= 2
 
-        # TODO(liym27): Uncomment code after "return" statement can be transformed correctly.
-        # if not self.is_train:
-        #     # get pred
-        #     yolo_boxes = fluid.layers.concat(self.boxes, axis=1)
-        #     yolo_scores = fluid.layers.concat(self.scores, axis=2)
-        #
-        #     pred = fluid.layers.multiclass_nms(
-        #         bboxes=yolo_boxes,
-        #         scores=yolo_scores,
-        #         score_threshold=cfg.valid_thresh,
-        #         nms_top_k=cfg.nms_topk,
-        #         keep_top_k=cfg.nms_posk,
-        #         nms_threshold=cfg.nms_thresh,
-        #         background_label=-1)
-        #     return pred
-        # else:
-        #     return sum(self.losses)
-        return sum(self.losses)
+        if not self.is_train:
+            # get pred
+            yolo_boxes = fluid.layers.concat(self.boxes, axis=1)
+            yolo_scores = fluid.layers.concat(self.scores, axis=2)
+
+            pred = fluid.layers.multiclass_nms(
+                bboxes=yolo_boxes,
+                scores=yolo_scores,
+                score_threshold=cfg.valid_thresh,
+                nms_top_k=cfg.nms_topk,
+                keep_top_k=cfg.nms_posk,
+                nms_threshold=cfg.nms_thresh,
+                background_label=-1)
+            return pred
+        else:
+            return sum(self.losses)