add --worker_num for yolov3 (#2955)

PaddlePaddle · Jul 29, 2019 · 0266b3c · 0266b3c
1 parent b0783c7
commit 0266b3c
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 4 deletions.
diff --git a/PaddleCV/yolov3/README.md b/PaddleCV/yolov3/README.md
@@ -95,6 +95,7 @@ dataset/coco/
 
 - 通过设置`export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7`指定8卡GPU训练。
 - 若在Windows环境下训练模型，建议设置`--use_multiprocess_reader=False`。
+- 通过`--worker_num=`设置多进程数据读取器进程数，默认进程数为8，若训练机器CPU核数较少，建议设小该值。
 - 可选参数见：
 
     python train.py --help

diff --git a/PaddleCV/yolov3/README_en.md b/PaddleCV/yolov3/README_en.md
@@ -96,6 +96,7 @@ Please make sure that pre-trained model is downloaded and loaded correctly, othe
 
 - Set `export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7` to specifiy 8 GPUs to train. 
 - It is recommended to set `--use_multiprocess_reader=False` when training on Windows.
+- Set `--worker_num=` to specifiy multiprocess reader worker number, which is default 8, if the number of CPU cores in the training environment is small, it is recommended to set worker number to a small value.
 - For more help on arguments:
 
     python train.py --help

diff --git a/PaddleCV/yolov3/reader.py b/PaddleCV/yolov3/reader.py
@@ -303,8 +303,9 @@ def train(size=416,
         return generator
     else:
         print("multiprocess is not fully compatible with Windows, "
-                "you can set --use_multiprocess_reader=False if there "
-                "are errors incured by multiprocess")
+                "you can set --use_multiprocess_reader=False if you "
+                "are training on Windows and there are errors incured "
+                "by multiprocess.")
         print("multiprocess reader starting up, it takes a while...")
 
     def infinite_reader():

diff --git a/PaddleCV/yolov3/train.py b/PaddleCV/yolov3/train.py
@@ -157,7 +157,8 @@ def if_exist(var):
         total_iter=total_iter * devices_num,
         mixup_iter=mixup_iter * devices_num,
         random_sizes=random_sizes,
-        use_multiprocess_reader=cfg.use_multiprocess_reader)
+        use_multiprocess_reader=cfg.use_multiprocess_reader,
+        num_workers=cfg.worker_num)
     py_reader = model.py_reader
     py_reader.decorate_paddle_reader(train_reader)
 

diff --git a/PaddleCV/yolov3/utility.py b/PaddleCV/yolov3/utility.py
@@ -120,7 +120,8 @@ def parse_args():
     add_arg('class_num',        int,    80,          "Class number.")
     add_arg('data_dir',         str,    'dataset/coco',        "The data root path.")
     add_arg('start_iter',       int,    0,      "Start iteration.")
-    add_arg('use_multiprocess_reader', bool,   True,   "add multiprocess.")
+    add_arg('use_multiprocess_reader', bool,   True,   "whether use multiprocess reader.")
+    add_arg('worker_num',       int,   8,   "worker number for multiprocess reader.")
     #SOLVER
     add_arg('batch_size',       int,    8,      "Mini-batch size per device.")
     add_arg('learning_rate',    float,  0.001,  "Learning rate.")