def test_save_best_after_loading_from_latest(self, temp_dir):
checkpointer = Checkpointer(temp_dir.path)
checkpointer.save_latest(self.state_dict)
checkpointer.load_model(self.model, self.optimizer, temp_dir.path,
"model.latest")
self.state_dict["score"] = 3.0
self.state_dict["epoch"] = 1
checkpointer.save_best(self.state_dict)
self.assertEqual(self.state_dict["score"], checkpointer.best_score)
self.assertEqual(self.state_dict["epoch"], 1)
def test_load_model_from_save_best(self, mock_torch_load, temp_dir):
self.state_dict["score"] = 5.0
mock_torch_load.return_value = self.state_dict
checkpointer = Checkpointer(temp_dir.path)
checkpointer.best_score = self.state_dict["score"] + 1
checkpointer.save_best(self.state_dict)
epoch_num, model, optimizer = checkpointer.load_model(
self.model, self.optimizer, temp_dir.path, "model.best")
self.assertTrue(mock_torch_load.call_count == 1)
self.check_loaded_objects(epoch_num, model, optimizer)
def test_save_best_folder_None(self, temp_dir):
checkpointer = Checkpointer(temp_dir.path)
scores = [3.0, 7.0, 2.0, 11.0]
best_epoch = [1, 1, 3, 3]
for epoch_num, score in enumerate(scores):
self.state_dict["epoch"] = epoch_num + 1
self.state_dict["score"] = score
with self.subTest(state_dict=self.state_dict,
epoch_num=epoch_num,
score=score,
best_score=checkpointer.best_score):
checkpointer.save_best(self.state_dict)
self.state_dict = torch.load(
os.path.join(temp_dir.path, "model.best"))
self.assertEqual(self.state_dict["score"],
checkpointer.best_score)
self.assertEqual(self.state_dict["epoch"],
best_epoch[epoch_num])