Add eval_batch_size (fix #17)

Author: arxyzan

audio/configs/wav2vec2-pretraining.yaml

@@ -15,7 +15,8 @@ dataset:
 optimizer:
   lr: 0.0005
 train:
-  batch_size: 1
+  batch_size: 16
+  eval_batch_size: 16
   num_epochs: 1000
   log_dir: 'audio/logs'
   save_ckpt_freq: 10

audio/trainer.py

@@ -36,7 +36,7 @@ def __init__(self, cfg: DictConfig):
                                                                 padding='longest')
         self.train_loader = DataLoader(self.train_dataset, batch_size=cfg.train.batch_size,
                                        collate_fn=self.data_collator)
-        self.test_loader = DataLoader(self.test_dataset, batch_size=cfg.train.val_batch_size,
+        self.test_loader = DataLoader(self.test_dataset, batch_size=cfg.train.eval_batch_size,
                                       collate_fn=self.data_collator)
         # Tensorboard
         self.tensorboard = SummaryWriter(log_dir=self.cfg.train.log_dir)

text/configs/roberta-pretraining.yaml

@@ -2,6 +2,7 @@ modality: 'text'
 device: 'cuda'
 train:
   batch_size: 32
+  eval_batch_size: 32
   num_epochs: 20
   checkpoints_dir: 'text/checkpoints/roberta-pretrain'
   log_dir: 'text/logs/roberta-pretrain'

text/trainer.py

@@ -44,7 +44,7 @@ def __init__(self, cfg: DictConfig):
         self.test_dataset = WikiText(cfg, 'test', self.tokenizer)
         self.train_loader = DataLoader(self.train_dataset, batch_size=cfg.train.batch_size,
                                        collate_fn=self.train_dataset.collate_fn)
-        self.test_loader = DataLoader(self.test_dataset, batch_size=cfg.train.val_batch_size,
+        self.test_loader = DataLoader(self.test_dataset, batch_size=cfg.train.eval_batch_size,
                                       collate_fn=self.test_dataset.collate_fn)
         # Tensorboard
         self.tensorboard = SummaryWriter(log_dir=self.cfg.train.log_dir)

vision/configs/beit-pretraining.yaml

@@ -25,6 +25,7 @@ dataset:
 train:
   num_epochs: 800
   batch_size: 16
+  eval_batch_size: 16
   shuffle: true
   save_ckpt_freq: 20
   checkpoints_dir: 'vision/checkpoints/beit-pretrain'

vision/trainer.py

@@ -31,7 +31,7 @@ def __init__(self, cfg):
         self.train_dataset = MIMPretrainingDataset(cfg, split='train')
         self.test_dataset = MIMPretrainingDataset(cfg, split='test')
         self.train_loader = DataLoader(self.train_dataset, batch_size=cfg.train.batch_size, shuffle=cfg.train.shuffle)
-        self.test_loader = DataLoader(self.test_dataset, batch_size=cfg.train.batch_size, shuffle=cfg.train.shuffle)
+        self.test_loader = DataLoader(self.test_dataset, batch_size=cfg.train.eval_batch_size, shuffle=cfg.train.shuffle)
 
         # Tensorboard
         self.tensorboard = SummaryWriter(log_dir=self.cfg.train.log_dir)