From b9d3e8eeec88d9522d2f391e4e50346e6e5cee39 Mon Sep 17 00:00:00 2001 From: WenmuZhou Date: Thu, 4 Feb 2021 11:33:48 +0800 Subject: [PATCH] add dataset len check --- tools/program.py | 5 +++++ tools/train.py | 5 +++++ 2 files changed, 10 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..a88d64c2 100755 --- a/tools/train.py +++ b/tools/train.py @@ -50,6 +50,11 @@ 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' + ) + if config['Eval']: valid_dataloader = build_dataloader(config, 'Eval', device, logger) else: -- GitLab