From 8447ec779b94d83e944db50c5af69b6491650230 Mon Sep 17 00:00:00 2001
From: wuzewu <wuzewu@baidu.com>
Date: Tue, 27 Aug 2019 11:40:29 +0800
Subject: [PATCH] Fix train log issue

---
 pdseg/train.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/pdseg/train.py b/pdseg/train.py
index 4ab3e968..05db0de8 100644
--- a/pdseg/train.py
+++ b/pdseg/train.py
@@ -192,12 +192,10 @@ def train(cfg):
 
     def data_generator():
         if args.use_mpio:
-            print("Use multiprocess reader")
             data_gen = dataset.multiprocess_generator(
                 num_processes=cfg.DATALOADER.NUM_WORKERS,
                 max_queue_size=cfg.DATALOADER.BUF_SIZE)
         else:
-            print("Use multi-thread reader")
             data_gen = dataset.generator()
 
         batch_data = []
@@ -332,6 +330,11 @@ def train(cfg):
             ("begin epoch[{}] is larger than cfg.SOLVER.NUM_EPOCHS[{}]").format(
                 begin_epoch, cfg.SOLVER.NUM_EPOCHS))
 
+    if args.use_mpio:
+        print("Use multiprocess reader")
+    else:
+        print("Use multi-thread reader")
+
     for epoch in range(begin_epoch, cfg.SOLVER.NUM_EPOCHS + 1):
         py_reader.start()
         while True:
-- 
GitLab