fix batch error for val

wjtscn · Dec 17, 2019 · def64e0 · def64e0
1 parent d84f6c3
commit def64e0
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/train_yolov3.py b/train_yolov3.py
@@ -203,7 +203,10 @@ def get_dataloader(net, train_dataset, val_dataset, batch_size):
             batch_size, True, batchify_fn=batchify_fn, last_batch='rollover', num_workers=FLAGS.num_workers)
 
         val_batchify_fn = Tuple(Stack(), Pad(pad_val=-1))
-        val_batchify_fn = Tuple(*([Stack() for _ in range(3)] + [Pad(axis=0, pad_val=-1) for _ in range(1)]))
+        if FLAGS.mult_out:
+            val_batchify_fn = Tuple(*([Stack() for _ in range(3)] + [Pad(axis=1, pad_val=-1) for _ in range(1)]))
+        else:
+            val_batchify_fn = Tuple(*([Stack() for _ in range(3)] + [Pad(axis=0, pad_val=-1) for _ in range(1)]))
         val_loader = gluon.data.DataLoader(
             val_dataset.transform(YOLO3NBVideoInferenceTransform(width, height)),
             batch_size, False, batchify_fn=val_batchify_fn, last_batch='discard', num_workers=FLAGS.num_workers)