break if your dataset size is smaller than your batch size

This commit is contained in:
mrq 2023-02-17 04:08:27 +00:00
parent 94d0f16608
commit 6afa2c299e

View File

@ -110,6 +110,8 @@ class Trainer:
self.dataset_debugger = get_dataset_debugger(dataset_opt)
if self.dataset_debugger is not None and resume_state is not None:
self.dataset_debugger.load_state(opt_get(resume_state, ['dataset_debugger_state'], {}))
if len(self.train_set) <= dataset_opt['batch_size']:
raise Exception("dataset size is less than batch size, consider reducing your batch size, or increasing your dataset.")
train_size = int(math.ceil(len(self.train_set) / dataset_opt['batch_size']))
total_iters = int(opt['train']['niter'])
self.total_epochs = int(math.ceil(total_iters / train_size))