From 8e9d8515631272d8990f7e5a72eb6a7d6659a294 Mon Sep 17 00:00:00 2001
From: WenmuZhou <zjwenmu@gmail.com>
Date: Fri, 5 Feb 2021 12:35:08 +0800
Subject: [PATCH] add dataset len check

---
 tools/program.py | 5 +++++
 tools/train.py   | 6 ++++++
 2 files changed, 11 insertions(+)

diff --git a/tools/program.py b/tools/program.py
index 99a37432..34d484d8 100755
--- a/tools/program.py
+++ b/tools/program.py
@@ -163,6 +163,11 @@ def train(config,
     if type(eval_batch_step) == list and len(eval_batch_step) >= 2:
         start_eval_step = eval_batch_step[0]
         eval_batch_step = eval_batch_step[1]
+        if len(valid_dataloader) == 0:
+            logger.info(
+                'No Images in eval dataset, evaluation during training will be disabled'
+            )
+            start_eval_step = 1e111
         logger.info(
             "During the training process, after the {}th iteration, an evaluation is run every {} iterations".
             format(start_eval_step, eval_batch_step))
diff --git a/tools/train.py b/tools/train.py
index 383f8d83..fab10b64 100755
--- a/tools/train.py
+++ b/tools/train.py
@@ -50,6 +50,12 @@ def main(config, device, logger, vdl_writer):
 
     # build dataloader
     train_dataloader = build_dataloader(config, 'Train', device, logger)
+    if len(train_dataloader) == 0:
+        logger.error(
+            'No Images in train dataset, please check annotation file and path in the configuration file'
+        )
+        return
+
     if config['Eval']:
         valid_dataloader = build_dataloader(config, 'Eval', device, logger)
     else:
-- 
GitLab