From 6afa2c299e4c4b25549c2c342f6e7e321344ef2a Mon Sep 17 00:00:00 2001 From: mrq Date: Fri, 17 Feb 2023 04:08:27 +0000 Subject: [PATCH] break if your dataset size is smaller than your batch size --- codes/train.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/codes/train.py b/codes/train.py index 690300f5..b9aa82c4 100644 --- a/codes/train.py +++ b/codes/train.py @@ -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))