add trainer debugging option "overfit_single_batch"

base/base_data_loader.py

@@ -8,14 +8,23 @@ class BaseDataLoader(DataLoader):
     """
     Base class for all data loaders
     """
-    def __init__(self, dataset, batch_size, shuffle, validation_split, num_workers, collate_fn=default_collate):
+    def __init__(self, dataset, batch_size, shuffle, validation_split, num_workers, collate_fn=default_collate,
+                 single_batch=False):
         self.validation_split = validation_split
         self.shuffle = shuffle
 
         self.batch_idx = 0
         self.n_samples = len(dataset)
 
-        self.sampler, self.valid_sampler = self._split_sampler(self.validation_split)
+        if not single_batch:
+            self.sampler, self.valid_sampler = self._split_sampler(self.validation_split)
+        else:
+            idx_full = np.arange(self.n_samples)
+            np.random.seed(0)
+            np.random.shuffle(idx_full)
+            self.sampler = SubsetRandomSampler(idx_full[:batch_size])
+            self.valid_sampler = None
+            self.shuffle = None
 
         self.init_kwargs = {
             'dataset': dataset,

base/base_trainer.py

@@ -26,6 +26,7 @@ def __init__(self, model, criterion, metric_ftns, optimizer, config):
         self.epochs = cfg_trainer['epochs']
         self.save_period = cfg_trainer['save_period']
         self.monitor = cfg_trainer.get('monitor', 'off')
+        self.overfit_single_batch = cfg_trainer.get('overfit_single_batch', False)
 
         # configuration to monitor model performance and save best
         if self.monitor == 'off':

config.json

@@ -45,6 +45,7 @@
         "monitor": "min val_loss",
         "early_stop": 10,
 
-        "tensorboard": true
+        "tensorboard": true,
+        "overfit_single_batch": false
     }
 }

train.py

@@ -17,11 +17,13 @@
 torch.backends.cudnn.benchmark = False
 np.random.seed(SEED)
 
+
 def main(config):
     logger = config.get_logger('train')
 
     # setup data_loader instances
-    data_loader = config.init_obj('data_loader', module_data)
+    data_loader = config.init_obj('data_loader', module_data,
+                                  single_batch=config['trainer'].get('overfit_single_batch', False))
     valid_data_loader = data_loader.split_validation()
 
     # build model architecture, then print to console

trainer/trainer.py

@@ -59,8 +59,9 @@ def _train_epoch(self, epoch):
                     loss.item()))
                 self.writer.add_image('input', make_grid(data.cpu(), nrow=8, normalize=True))
 
-            if batch_idx == self.len_epoch:
+            if batch_idx == self.len_epoch or self.overfit_single_batch:
                 break
+
         log = self.train_metrics.result()
 
         if self.do_validation:
@@ -81,7 +82,7 @@ def _valid_epoch(self, epoch):
         self.model.eval()
         self.valid_metrics.reset()
         with torch.no_grad():
-            for batch_idx, (data, target) in enumerate(self.valid_data_loader):
+            for batch_idx, (data, target, _) in enumerate(self.valid_data_loader):
                 data, target = data.to(self.device), target.to(self.device)
 
                 output = self.model(data)