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)
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_pretrained_encoder_with_dict_attr(self, temp_dir):
checkpointer = Checkpointer(temp_dir.path)
checkpointer.save_latest(self.state_dict, filename=self.model_name)
encoder = PretrainedEncoder(FullyConnectedMapper,
pretrained_path=os.path.join(
temp_dir.path, self.model_name),
encoder_args=(2, 3),
checkpoint_key='model')
encoded = encoder(self.input)
expected_encoding = self.model(self.input)
self.assertEqual((encoded - expected_encoding).sum().data.numpy(), 0)
def test_save_latest_folder_None(self, temp_dir):
checkpointer = Checkpointer(temp_dir.path)
scores = [3.0, 7.0, 2.0, 11.0]
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):
checkpointer.save_latest(self.state_dict)
self.state_dict = torch.load(
os.path.join(temp_dir.path, "model.latest"))
self.assertEqual(self.state_dict["score"], score)
self.assertEqual(self.state_dict["epoch"], epoch_num + 1)