def test_load_checkpoint_user_dict(aggregate_checkpoints_mock, load_mock):
trainer = _create_trainer()
trainer_options = {
'mixed_precision': np.bool_(False),
'world_rank': np.int64(0),
'world_size': np.int64(1),
'horizontal_parallel_size': np.int64(1),
'data_parallel_size': np.int64(1),
'zero_stage': np.int64(0)
}
state_dict = {
'model': {},
'optimizer': {},
'trainer_options': {
'mixed_precision': np.bool_(False),
'world_rank': np.int64(0),
'world_size': np.int64(1),
'horizontal_parallel_size': np.int64(1),
'data_parallel_size': np.int64(1),
'zero_stage': np.int64(0)
},
'user_dict':
_checkpoint_storage.to_serialized_hex(
{'array': torch.tensor(np.arange(5))})
}
trainer.load_state_dict = Mock()
load_mock.side_effect = [trainer_options, state_dict]
user_dict = trainer.load_checkpoint('abc')
assert torch.all(torch.eq(user_dict['array'], torch.tensor(np.arange(5))))
def test_load_checkpoint_user_dict(aggregate_checkpoints_mock, load_mock):
trainer = _create_trainer()
trainer_options = {
"mixed_precision": np.bool_(False),
"world_rank": np.int64(0),
"world_size": np.int64(1),
"horizontal_parallel_size": np.int64(1),
"data_parallel_size": np.int64(1),
"zero_stage": np.int64(0),
}
state_dict = {
"model": {},
"optimizer": {},
"trainer_options": {
"mixed_precision": np.bool_(False),
"world_rank": np.int64(0),
"world_size": np.int64(1),
"horizontal_parallel_size": np.int64(1),
"data_parallel_size": np.int64(1),
"zero_stage": np.int64(0),
},
"user_dict":
_checkpoint_storage.to_serialized_hex(
{"array": torch.tensor(np.arange(5))}),
}
trainer.load_state_dict = Mock()
load_mock.side_effect = [trainer_options, state_dict]
user_dict = trainer.load_checkpoint("abc")
assert torch.all(torch.eq(user_dict["array"], torch.tensor(np.arange(5))))
def test_save_checkpoint_user_dict(save_mock):
trainer = _create_trainer()
state_dict = {'model': {}, 'optimizer': {}}
trainer.state_dict = Mock(return_value=state_dict)
trainer.save_checkpoint('abc', user_dict={'abc': np.arange(4)})
save_args, _ = save_mock.call_args
assert 'user_dict' in save_args[0]
assert save_args[0]['user_dict'] == _checkpoint_storage.to_serialized_hex(
{'abc': np.arange(4)})
def test_save_checkpoint_user_dict(save_mock):
trainer = _create_trainer()
state_dict = {"model": {}, "optimizer": {}}
trainer.state_dict = Mock(return_value=state_dict)
trainer.save_checkpoint("abc", user_dict={"abc": np.arange(4)})
save_args, _ = save_mock.call_args
assert "user_dict" in save_args[0]
assert save_args[0]["user_dict"] == _checkpoint_storage.to_serialized_hex(
{"abc": np.arange(4)})