def test_trainer_respects_epoch_size_smaller_tnan_total(self):
batches_per_epoch = 1
num_epochs = 2
data_loader_smaller_epoch = AllennlpDataLoader(
self.instances,
batch_size=2,
collate_fn=allennlp_collate,
batches_per_epoch=batches_per_epoch,
)
trainer = GradientDescentTrainer(
self.model,
self.optimizer,
data_loader_smaller_epoch,
validation_data_loader=self.validation_data_loader,
num_epochs=num_epochs,
serialization_dir=self.TEST_DIR,
)
assert trainer._batch_num_total == 0
metrics = trainer.train()
epoch = metrics["epoch"]
assert epoch == num_epochs - 1
assert trainer._batch_num_total == num_epochs * batches_per_epoch
def test_data_loader_lazy_epoch_size_correct_custom_epoch_size(self):
batches_per_epoch = 3
num_epochs = 3
data_loader_custom_epoch_lazy = AllennlpDataLoader(
self.instances_lazy,
batch_size=2,
collate_fn=allennlp_collate,
batches_per_epoch=batches_per_epoch,
)
trainer = GradientDescentTrainer(
self.model,
self.optimizer,
data_loader_custom_epoch_lazy,
validation_data_loader=self.validation_data_loader,
num_epochs=num_epochs,
serialization_dir=self.TEST_DIR,
)
assert trainer._batch_num_total == 0
metrics = trainer.train()
epoch = metrics["epoch"]
assert epoch == num_epochs - 1
assert trainer._batch_num_total == num_epochs * batches_per_epoch