From db216ad20a24adb5efe0a9c5aed2b8c698de58d8 Mon Sep 17 00:00:00 2001
From: hysunflower <52739577+hysunflower@users.noreply.github.com>
Date: Tue, 12 Nov 2019 15:54:19 +0800
Subject: [PATCH] add profiler and max_iter for yolov3 model (#3908)

---
 PaddleCV/yolov3/train.py   | 8 ++++++++
 PaddleCV/yolov3/utility.py | 3 +++
 2 files changed, 11 insertions(+)

diff --git a/PaddleCV/yolov3/train.py b/PaddleCV/yolov3/train.py
index 5f2284cf..a89c6b30 100644
--- a/PaddleCV/yolov3/train.py
+++ b/PaddleCV/yolov3/train.py
@@ -41,6 +41,7 @@ from utility import (parse_args, print_arguments,
 
 import paddle
 import paddle.fluid as fluid
+from paddle.fluid import profiler
 import reader
 from models.yolov3 import YOLOv3
 from learning_rate import exponential_with_warmup_decay
@@ -186,6 +187,13 @@ def train():
                 iter_id, lr[0],
                 smoothed_loss.get_mean_value(), start_time - prev_start_time))
             sys.stdout.flush()
+            #add profiler tools
+            if args.is_profiler and iter_id == 5:
+               profiler.start_profiler("All")
+            elif args.is_profiler and iter_id == 10:
+                 profiler.stop_profiler("total", args.profiler_path)
+                 return
+
             if (iter_id + 1) % cfg.snapshot_iter == 0:
                 save_model("model_iter{}".format(iter_id))
                 print("Snapshot {} saved, average loss: {}, \
diff --git a/PaddleCV/yolov3/utility.py b/PaddleCV/yolov3/utility.py
index 3e0f3082..9d442f4e 100644
--- a/PaddleCV/yolov3/utility.py
+++ b/PaddleCV/yolov3/utility.py
@@ -146,6 +146,9 @@ def parse_args():
     add_arg('draw_thresh',      float,  0.5,
             "Confidence score threshold to draw prediction box in image in debug mode")
     add_arg('enable_ce',        bool,  False,                "If set True, enable continuous evaluation job.")
+    # args for profiler tools
+    add_arg('is_profiler',        int,  0,                "the switch of profiler")
+    add_arg('profiler_path',        str,  './',                "the path to save profiler output files")
     # yapf: enable
     args = parser.parse_args()
     file_name = sys.argv[0]
-- 
GitLab