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_predefined_folder(self, mock_torch_load,
temp_dir):
checkpointer = Checkpointer(temp_dir.path)
epoch_num, model, optimizer = checkpointer.load_model(
self.model, self.optimizer, temp_dir.path)
mock_torch_load.assert_not_called()
self.check_loaded_objects(epoch_num, model, optimizer)
def test_load_model_from_save_latest(self, mock_torch_load, temp_dir):
mock_torch_load.return_value = self.state_dict
checkpointer = Checkpointer(temp_dir.path)
checkpointer.save_latest(self.state_dict)
epoch_num, model, optimizer = checkpointer.load_model(
self.model, self.optimizer, temp_dir.path, "model.latest")
self.assertTrue(mock_torch_load.call_count == 1)
self.check_loaded_objects(epoch_num, model, optimizer)