From 43122c2027ecb7a420e25e9fbd38c875a8495e87 Mon Sep 17 00:00:00 2001 From: Dang Qingqing Date: Thu, 12 Apr 2018 22:30:27 +0800 Subject: [PATCH] Fix reader bug. --- fluid/object_detection/train.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/fluid/object_detection/train.py b/fluid/object_detection/train.py index 0f2856ca..09a70aed 100644 --- a/fluid/object_detection/train.py +++ b/fluid/object_detection/train.py @@ -174,6 +174,9 @@ def parallel_exe(args, elif data_args.dataset == 'pascalvoc': num_classes = 21 + devices = os.getenv("CUDA_VISIBLE_DEVICES") or "" + devices_num = len(devices.split(",")) + image = fluid.layers.data(name='image', shape=image_shape, dtype='float32') gt_box = fluid.layers.data( name='gt_box', shape=[4], dtype='float32', lod_level=1) @@ -253,6 +256,7 @@ def parallel_exe(args, for batch_id, data in enumerate(train_reader()): prev_start_time = start_time start_time = time.time() + if len(data) < devices_num: continue loss_v, = train_exe.run(fetch_list=[loss.name], feed_dict=feeder.feed(data)) end_time = time.time() -- GitLab