add dataset len check

This commit is contained in:
WenmuZhou 2021-02-05 12:35:08 +08:00
parent 60effbd748
commit 8e9d851563
2 changed files with 11 additions and 0 deletions

View File

@ -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))

View File

@ -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: