def create_model_saver( trainer_settings: TrainerSettings, model_path: str, load: bool ) -> BaseModelSaver: model_saver = TorchModelSaver( # type: ignore trainer_settings, model_path, load ) return model_saver
def create_model_saver(framework: str, trainer_settings: TrainerSettings, model_path: str, load: bool) -> BaseModelSaver: if framework == FrameworkType.PYTORCH: model_saver = TorchModelSaver( # type: ignore trainer_settings, model_path, load) else: model_saver = TFModelSaver( # type: ignore trainer_settings, model_path, load) return model_saver
def test_checkpoint_conversion(tmpdir, rnn, visual, discrete): dummy_config = TrainerSettings() model_path = os.path.join(tmpdir, "Mock_Brain") policy = create_policy_mock( dummy_config, use_rnn=rnn, use_discrete=discrete, use_visual=visual ) trainer_params = TrainerSettings() model_saver = TorchModelSaver(trainer_params, model_path) model_saver.register(policy) model_saver.save_checkpoint("Mock_Brain", 100) assert os.path.isfile(model_path + "/Mock_Brain-100.onnx")
def test_register(tmp_path): trainer_params = TrainerSettings() model_saver = TorchModelSaver(trainer_params, tmp_path) opt = mock.Mock(spec=TorchPPOOptimizer) opt.get_modules = mock.Mock(return_value={}) model_saver.register(opt) assert model_saver.policy is None trainer_params = TrainerSettings() policy = create_policy_mock(trainer_params) opt.get_modules = mock.Mock(return_value={}) model_saver.register(policy) assert model_saver.policy is not None
def test_load_policy_different_hidden_units(tmp_path, vis_encode_type): path1 = os.path.join(tmp_path, "runid1") trainer_params = TrainerSettings() trainer_params.network_settings = NetworkSettings( hidden_units=12, vis_encode_type=EncoderType(vis_encode_type)) policy = create_policy_mock(trainer_params, use_visual=True) conv_params = [ mod for mod in policy.actor.parameters() if len(mod.shape) > 2 ] model_saver = TorchModelSaver(trainer_params, path1) model_saver.register(policy) model_saver.initialize_or_load(policy) policy.set_step(2000) mock_brain_name = "MockBrain" model_saver.save_checkpoint(mock_brain_name, 2000) # Try load from this path trainer_params2 = TrainerSettings() trainer_params2.network_settings = NetworkSettings( hidden_units=10, vis_encode_type=EncoderType(vis_encode_type)) model_saver2 = TorchModelSaver(trainer_params2, path1, load=True) policy2 = create_policy_mock(trainer_params2, use_visual=True) conv_params2 = [ mod for mod in policy2.actor.parameters() if len(mod.shape) > 2 ] # asserts convolutions have different parameters before load for conv1, conv2 in zip(conv_params, conv_params2): assert not torch.equal(conv1, conv2) # asserts layers still have different dimensions for mod1, mod2 in zip(policy.actor.parameters(), policy2.actor.parameters()): if mod1.shape[0] == 12: assert mod2.shape[0] == 10 model_saver2.register(policy2) model_saver2.initialize_or_load(policy2) # asserts convolutions have same parameters after load for conv1, conv2 in zip(conv_params, conv_params2): assert torch.equal(conv1, conv2) # asserts layers still have different dimensions for mod1, mod2 in zip(policy.actor.parameters(), policy2.actor.parameters()): if mod1.shape[0] == 12: assert mod2.shape[0] == 10
def test_load_save_policy(tmp_path): path1 = os.path.join(tmp_path, "runid1") path2 = os.path.join(tmp_path, "runid2") trainer_params = TrainerSettings() policy = create_policy_mock(trainer_params) model_saver = TorchModelSaver(trainer_params, path1) model_saver.register(policy) model_saver.initialize_or_load(policy) policy.set_step(2000) mock_brain_name = "MockBrain" model_saver.save_checkpoint(mock_brain_name, 2000) assert len(os.listdir(tmp_path)) > 0 # Try load from this path model_saver2 = TorchModelSaver(trainer_params, path1, load=True) policy2 = create_policy_mock(trainer_params) model_saver2.register(policy2) model_saver2.initialize_or_load(policy2) _compare_two_policies(policy, policy2) assert policy2.get_current_step() == 2000 # Try initialize from path 1 trainer_params.init_path = path1 model_saver3 = TorchModelSaver(trainer_params, path2) policy3 = create_policy_mock(trainer_params) model_saver3.register(policy3) model_saver3.initialize_or_load(policy3) _compare_two_policies(policy2, policy3) # Assert that the steps are 0. assert policy3.get_current_step() == 0
def test_load_save_optimizer(tmp_path, optimizer): OptimizerClass, HyperparametersClass = optimizer trainer_settings = TrainerSettings() trainer_settings.hyperparameters = HyperparametersClass() policy = create_policy_mock(trainer_settings, use_discrete=False) optimizer = OptimizerClass(policy, trainer_settings) # save at path 1 path1 = os.path.join(tmp_path, "runid1") model_saver = TorchModelSaver(trainer_settings, path1) model_saver.register(policy) model_saver.register(optimizer) model_saver.initialize_or_load() policy.set_step(2000) model_saver.save_checkpoint("MockBrain", 2000) # create a new optimizer and policy policy2 = create_policy_mock(trainer_settings, use_discrete=False) optimizer2 = OptimizerClass(policy2, trainer_settings) # load weights model_saver2 = TorchModelSaver(trainer_settings, path1, load=True) model_saver2.register(policy2) model_saver2.register(optimizer2) model_saver2.initialize_or_load() # This is to load the optimizers # Compare the two optimizers _compare_two_optimizers(optimizer, optimizer2)
def test_reward_provider_save(tmp_path, optimizer): OptimizerClass, HyperparametersClass = optimizer trainer_settings = TrainerSettings() trainer_settings.hyperparameters = HyperparametersClass() trainer_settings.reward_signals = { RewardSignalType.CURIOSITY: CuriositySettings(), RewardSignalType.GAIL: GAILSettings(demo_path=DEMO_PATH), } policy = create_policy_mock(trainer_settings, use_discrete=False) optimizer = OptimizerClass(policy, trainer_settings) # save at path 1 path1 = os.path.join(tmp_path, "runid1") model_saver = TorchModelSaver(trainer_settings, path1) model_saver.register(policy) model_saver.register(optimizer) model_saver.initialize_or_load() policy.set_step(2000) model_saver.save_checkpoint("MockBrain", 2000) # create a new optimizer and policy optimizer2 = OptimizerClass(policy, trainer_settings) policy2 = create_policy_mock(trainer_settings, use_discrete=False) # load weights model_saver2 = TorchModelSaver(trainer_settings, path1, load=True) model_saver2.register(policy2) model_saver2.register(optimizer2) model_saver2.initialize_or_load() # This is to load the optimizers # assert the models have the same weights module_dict_1 = optimizer.get_modules() module_dict_2 = optimizer2.get_modules() assert "Module:GAIL" in module_dict_1 assert "Module:GAIL" in module_dict_2 for name, module1 in module_dict_1.items(): assert name in module_dict_2 module2 = module_dict_2[name] if hasattr(module1, "parameters"): for param1, param2 in zip(module1.parameters(), module2.parameters()): assert param1.data.ne(param2.data).sum() == 0
def test_load_different_reward_provider(caplog, tmp_path, optimizer): OptimizerClass, HyperparametersClass = optimizer trainer_settings = TrainerSettings() trainer_settings.hyperparameters = HyperparametersClass() trainer_settings.reward_signals = { RewardSignalType.CURIOSITY: CuriositySettings(), RewardSignalType.RND: RNDSettings(), } policy = create_policy_mock(trainer_settings, use_discrete=False) optimizer = OptimizerClass(policy, trainer_settings) # save at path 1 path1 = os.path.join(tmp_path, "runid1") model_saver = TorchModelSaver(trainer_settings, path1) model_saver.register(policy) model_saver.register(optimizer) model_saver.initialize_or_load() assert len(optimizer.critic.value_heads.stream_names) == 2 policy.set_step(2000) model_saver.save_checkpoint("MockBrain", 2000) trainer_settings2 = TrainerSettings() trainer_settings2.hyperparameters = HyperparametersClass() trainer_settings2.reward_signals = { RewardSignalType.GAIL: GAILSettings(demo_path=DEMO_PATH) } # create a new optimizer and policy policy2 = create_policy_mock(trainer_settings2, use_discrete=False) optimizer2 = OptimizerClass(policy2, trainer_settings2) # load weights model_saver2 = TorchModelSaver(trainer_settings2, path1, load=True) model_saver2.register(policy2) model_saver2.register(optimizer2) assert len(optimizer2.critic.value_heads.stream_names) == 1 model_saver2.initialize_or_load() # This is to load the optimizers messages = [ rec.message for rec in caplog.records if rec.levelno == WARNING ] assert len(messages) > 0