def test_keep_zero(self): checkpointer = Checkpointer(serialization_dir=self.TEST_DIR, keep_most_recent_by_count=0) for epochs_completed in range(5): state = { "epochs_completed": epochs_completed, "batches_in_epoch_completed": 0 } checkpointer.maybe_save_checkpoint( FakeTrainer(model_state=state, training_state=state), epochs_completed, 0) files = os.listdir(self.TEST_DIR) assert not any("model_state_" in x for x in files) assert not any("training_state_" in x for x in files)
def test_with_time(self): num_epochs = 30 pauses = [5, 18, 26] target = [(e, 0) for e in pauses] checkpointer = Checkpointer( serialization_dir=self.TEST_DIR, save_completed_epochs=False, save_every_num_seconds=1, keep_most_recent_by_count=3, ) for e in range(num_epochs): if e in pauses: time.sleep(2) state = {"epochs_completed": e, "batches_in_epoch_completed": 0} checkpointer.maybe_save_checkpoint( trainer=FakeTrainer(model_state=state, training_state=state), num_epochs_completed=e, num_batches_in_epoch_completed=0, ) models, training = self.retrieve_and_delete_saved() assert models == training == target
def test_default(self): """ Tests that the default behavior keeps just the last 2 checkpoints. """ default_num_to_keep = 2 num_epochs = 5 target = [(e, 0) for e in range(num_epochs - default_num_to_keep, num_epochs)] checkpointer = Checkpointer(serialization_dir=self.TEST_DIR) for epochs_completed in range(num_epochs): for batches_completed in [0, 5, 10]: state = { "epochs_completed": epochs_completed, "batches_in_epoch_completed": batches_completed, } checkpointer.maybe_save_checkpoint( FakeTrainer(model_state=state, training_state=state), epochs_completed, batches_completed, ) models, training = self.retrieve_and_delete_saved() assert models == training == target