def test_unflatten_params(self): for module_init_fn in self._get_module_init_fns(): module = FlattenParamsWrapper(module_init_fn()) buffers = {k.replace("_fpw_module.", "") for k, _ in module.named_buffers()} def clone_state_dict(): return OrderedDict((k, v.clone()) for k, v in module.state_dict().items()) ref_flat_param = module.flat_param.clone() with module.unflatten_params(): ref_state_dict = clone_state_dict() assert not torch.all(ref_flat_param == 0) # confirm that unflatten_params reflects values from new_flat_param new_flat_param = torch.full_like(module.flat_param, fill_value=42.0) with module.unflatten_params(flat_param=new_flat_param): new_state_dict = clone_state_dict() assert new_state_dict.keys() == ref_state_dict.keys() for k, v in new_state_dict.items(): if k in buffers: # buffers are not changed torch.testing.assert_allclose(v, ref_state_dict[k]) else: # params reflect new_flat_param value assert torch.all(v == 42.0) # after context manager exits, we go back to previous (reference) state torch.testing.assert_allclose(module.flat_param, ref_flat_param) with module.unflatten_params(): ref_state_dict2 = clone_state_dict() assert objects_are_equal(ref_state_dict, ref_state_dict2) # if we load the new_state_dict, then the flat param should match new_flat_param module.load_state_dict(new_state_dict) torch.testing.assert_allclose(module.flat_param, new_flat_param)
def test_load_state_dict(self): module = self._get_shared_params_transformer() ref_state_dict = module.state_dict() ref_output = self._get_output(module) module = self._get_shared_params_transformer(seed=1234) flat_module = FlattenParamsWrapper(module) flat_module.load_state_dict(ref_state_dict) flat_output = self._get_output(flat_module) assert objects_are_equal(ref_output, flat_output)
def test_flat_state_dict(self): """Test that flat state dict can be reloaded and produces the same results.""" for module_init_fn in self._get_module_init_fns(): flat_module = FlattenParamsWrapper(module_init_fn()) ref_output = self._get_output(flat_module) flat_state_dict = flat_module.flat_state_dict() new_module = FlattenParamsWrapper(module_init_fn(seed=1234)) new_module.load_state_dict(flat_state_dict) new_output = self._get_output(new_module) assert objects_are_equal(ref_output, new_output)
def test_flat_state_dict(self): flat_module = self._get_shared_params_transformer() flat_module = FlattenParamsWrapper(flat_module) ref_output = self._get_output(flat_module) flat_state_dict = flat_module.flat_state_dict() new_module = self._get_shared_params_transformer(seed=1234) new_module = FlattenParamsWrapper(new_module) new_module.load_state_dict(flat_state_dict) new_output = self._get_output(new_module) assert objects_are_equal(ref_output, new_output)
def test_load_state_dict(self): """Test that original (unwrapped) state_dict can be loaded in wrapped module.""" for module_init_fn in self._get_module_init_fns(): module = module_init_fn() ref_state_dict = module.state_dict() ref_output = self._get_output(module) module = module_init_fn(seed=1234) flat_module = FlattenParamsWrapper(module) # This should work without the unflatten_params context manager flat_module.load_state_dict(ref_state_dict) flat_output = self._get_output(flat_module) assert objects_are_equal(ref_output, flat_output) # And it should work with the context manager too with flat_module.unflatten_params(): flat_module.load_state_dict(ref_state_dict) flat_output = self._get_output(flat_module) assert objects_are_equal(ref_output, flat_output)