| # Owner(s): ["module: nn"] | |
| import unittest | |
| import sys | |
| import os | |
| import subprocess | |
| import torch | |
| import torch.nn.utils.stateless as stateless | |
| from torch.testing._internal.common_cuda import TEST_MULTIGPU | |
| from torch.testing._internal.common_utils import run_tests, TestCase | |
| class MockModule(torch.nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.l1 = torch.nn.Linear(1, 1) | |
| self.register_buffer('buffer', torch.ones(1)) | |
| def forward(self, x): | |
| return self.l1(x) + self.buffer | |
| class TestStatelessFunctionalAPI(TestCase): | |
| def _run_call_with_mock_module(self, module, device='cpu', prefix=''): | |
| x = torch.rand((1, 1)).to(device) | |
| weight = torch.tensor([[1.0]], device=device) | |
| bias = torch.tensor([0.0], device=device) | |
| buffer = torch.tensor([0.0], device=device) | |
| if prefix != '': | |
| parameters = {f'{prefix}.l1.weight': weight, | |
| f'{prefix}.l1.bias': bias, | |
| f'{prefix}.buffer': buffer} | |
| else: | |
| parameters = {'l1.weight': weight, | |
| 'l1.bias': bias, | |
| 'buffer': buffer} | |
| to_check = module | |
| if prefix != '': | |
| to_check = getattr(module, prefix) | |
| prev_weight = to_check.l1.weight.clone() | |
| prev_buffer = to_check.buffer.clone() | |
| # the parameters represent an identity function contrary to the | |
| # existing params in module. So here we expect the result to be the | |
| # same as the input if the weight swapping went well. | |
| res = stateless.functional_call(module, parameters, x) | |
| self.assertEqual(x, res) | |
| # check that the weight remain unmodified | |
| cur_weight = to_check.l1.weight | |
| cur_buffer = to_check.buffer | |
| self.assertEqual(cur_weight, prev_weight) | |
| self.assertEqual(cur_buffer, prev_buffer) | |
| def test_functional_call(self): | |
| module = MockModule() | |
| self._run_call_with_mock_module(module) | |
| def test_functional_call_with_jit(self): | |
| module = MockModule() | |
| jit_module = torch.jit.script(module) | |
| with self.assertRaisesRegex( | |
| RuntimeError, | |
| r'used with Jitted modules' | |
| ): | |
| self._run_call_with_mock_module(jit_module) | |
| x = torch.rand((1, 1)) | |
| traced_module = torch.jit.trace(module, x) | |
| with self.assertRaisesRegex( | |
| RuntimeError, | |
| r'used with Jitted modules' | |
| ): | |
| self._run_call_with_mock_module(traced_module) | |
| @unittest.skipIf(not TEST_MULTIGPU, 'multi-GPU not supported') | |
| @unittest.skip("This doesn't work right now") | |
| def test_functional_call_with_data_parallel(self): | |
| module = MockModule() | |
| module.cuda() | |
| dp_module = torch.nn.DataParallel(module, [0, 1]) | |
| self._run_call_with_mock_module(dp_module, device='cuda', prefix='module') | |
| def test_functional_call_with_gradient(self): | |
| module = MockModule() | |
| x = torch.rand((1, 1)) | |
| weight = torch.tensor([[1.0]], requires_grad=True) | |
| bias = torch.tensor([0.0], requires_grad=True) | |
| buffer = torch.tensor([0.0]) | |
| parameters = {'l1.weight': weight, | |
| 'l1.bias': bias, | |
| 'buffer': buffer} | |
| res = stateless.functional_call(module, parameters, x) | |
| # Check that a backward step calculates the gradient of the supplied parameters | |
| res.backward() | |
| self.assertIsNotNone(weight.grad) | |
| self.assertIsNotNone(bias.grad) | |
| self.assertIsNone(buffer.grad) | |
| # Gradient was not calculated for the module stated and buffers | |
| self.assertIsNone(module.l1.weight.grad) | |
| self.assertIsNone(module.l1.bias.grad) | |
| self.assertIsNone(module.buffer.grad) | |
| def test_functional_batch_norm(self): | |
| module = torch.nn.BatchNorm1d(10) | |
| module.train() # Allow stats update | |
| # lets replace the running_mean buffer and check if its correctly updated | |
| x = torch.full((20, 10), 128.0) | |
| rm = torch.zeros(10) | |
| parameters = {'running_mean': rm} | |
| prev_rm = module.running_mean.clone() | |
| res = stateless.functional_call(module, parameters, x) | |
| cur_rm = module.running_mean | |
| self.assertEqual(cur_rm, prev_rm) | |
| self.assertEqual(rm, torch.full((10,), 12.8)) | |
| # Now run functional without reparametrization and check that the module has | |
| # been updated | |
| res = stateless.functional_call(module, {}, x) | |
| self.assertEqual(module.running_mean, torch.full((10,), 12.8)) | |
| def test_circular_references(self): | |
| module = MockModule() | |
| # Add a circular reference | |
| module.l1.m = module | |
| x = torch.rand((1, 1)) | |
| weight = torch.tensor([[1.0]]) | |
| bias = torch.tensor([0.0]) | |
| buffer = torch.tensor([0.0]) | |
| parameters = {'l1.m.l1.weight': weight, | |
| 'l1.bias': bias, | |
| 'l1.m.buffer': buffer} | |
| prev_weight = module.l1.weight.clone() | |
| prev_buffer = module.buffer.clone() | |
| res = stateless.functional_call(module, parameters, x) | |
| self.assertEqual(x, res) | |
| # check that the weights remain unmodified and were correctly accesed | |
| cur_weight = module.l1.weight | |
| cur_buffer = module.buffer | |
| self.assertEqual(cur_weight, prev_weight) | |
| self.assertEqual(cur_buffer, prev_buffer) | |
| def test_reparametrized_module_change_parametrization_original(self): | |
| module = MockModule() | |
| torch.nn.utils.parametrizations.spectral_norm(module.l1) | |
| self.assertTrue('l1.parametrizations.weight.original' in dict(module.named_parameters())) | |
| orig_sn_weight = module.l1.weight.clone() | |
| x = torch.rand((1, 1)) | |
| # We substitute the parameter inside the parametrization | |
| # the parametrization itself is not overwritten so it will be applied with a different | |
| # value for the original tensor | |
| parameters = {'l1.parametrizations.weight.original': torch.nn.Parameter(torch.tensor([[1.0]])), | |
| 'l1.bias': torch.tensor([0.0]), | |
| 'buffer': torch.tensor([0.0])} | |
| res = stateless.functional_call(module, parameters, x) | |
| self.assertEqual(x, res) | |
| # verify that the spectral normalization is still applied | |
| self.assertTrue('l1.parametrizations.weight.original' in dict(module.named_parameters())) | |
| self.assertEqual(orig_sn_weight, module.l1.weight) | |
| def test_reparamertize_module_fail_reset_to_original(self): | |
| module = MockModule() | |
| torch.nn.utils.parametrizations.spectral_norm(module.l1) | |
| self.assertTrue('l1.parametrizations.weight.original' in dict(module.named_parameters())) | |
| orig_sn_weight = module.l1.weight.clone() | |
| # We substitute the parameter inside the parametrization | |
| # the parametrization itself is not overwritten so it will be applied with a different | |
| # value for the original tensor | |
| parameters = {'l1.parametrizations.weight.original': torch.nn.Parameter(torch.tensor([[1.0]])), | |
| 'l1.bias': torch.tensor([0.0]), | |
| 'buffer': torch.tensor([0.0])} | |
| with self.assertRaisesRegex(RuntimeError, "shapes cannot be multiplied"): | |
| x = torch.rand((4, 5)) # to work, it should be of size (1, 1) | |
| stateless.functional_call(module, parameters, x) # this call will fail because x is the wrong size | |
| # verify that the spectral normalization is still applied | |
| self.assertTrue('l1.parametrizations.weight.original' in dict(module.named_parameters())) | |
| self.assertEqual(orig_sn_weight, module.l1.weight) | |
| def test_tied_weights_warns(self): | |
| module = MockModule() | |
| module.tied_bias = module.l1.bias | |
| module.register_buffer("tied_buffer", module.buffer) | |
| weight = torch.tensor([[1.0]],) | |
| bias = torch.tensor([0.0]) | |
| buffer = torch.tensor([0.0]) | |
| parameters = {'l1.weight': weight, | |
| 'l1.bias': bias, | |
| 'buffer': buffer} | |
| x = torch.randn(1, 1) | |
| self.assertNotWarn(lambda: stateless.functional_call(module, parameters, x)) | |
| # if tied values are the same tensors, shouldn't warn | |
| parameters['tied_bias'] = bias | |
| parameters['tied_buffer'] = buffer | |
| self.assertNotWarn(lambda: stateless.functional_call(module, parameters, x)) | |
| del parameters['tied_bias'] | |
| del parameters['tied_buffer'] | |
| with self.assertWarnsOnceRegex(UserWarning, "functional_call was passed multiple values"): | |
| parameters['tied_bias'] = torch.tensor([5.0]) | |
| stateless.functional_call(module, parameters, x) | |
| del parameters['tied_bias'] | |
| with self.assertWarnsOnceRegex(UserWarning, "functional_call was passed multiple values"): | |
| parameters['tied_buffer'] = torch.tensor([5.0]) | |
| stateless.functional_call(module, parameters, x) | |
| def test_setattr(self): | |
| class Foo(torch.nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.register_buffer('foo', torch.zeros(())) | |
| def forward(self, x): | |
| self.foo = self.foo + 1 | |
| return x + self.foo | |
| a = {'foo': torch.zeros(())} | |
| mod = Foo() | |
| stateless.functional_call(mod, a, torch.ones(())) | |
| self.assertEqual(mod.foo, torch.zeros(())) | |
| self.assertEqual(a['foo'], torch.ones(())) | |
| class TestStatelessDeprecation(TestCase): | |
| def test_private_stateless_warns(self): | |
| script = """ | |
| import torch | |
| import warnings | |
| with warnings.catch_warnings(record=True) as w: | |
| from torch.nn.utils import _stateless | |
| exit(len(w)) | |
| """ | |
| try: | |
| subprocess.check_output( | |
| [sys.executable, '-W', 'all', '-c', script], | |
| stderr=subprocess.STDOUT, | |
| # On Windows, opening the subprocess with the default CWD makes `import torch` | |
| # fail, so just set CWD to this script's directory | |
| cwd=os.path.dirname(os.path.realpath(__file__)),) | |
| except subprocess.CalledProcessError as e: | |
| self.assertEqual(e.returncode, 1) | |
| else: | |
| self.assertTrue(False, "No warning was raised.") | |
| class TestPythonOptimizeMode(TestCase): | |
| def test_runs_with_optimize_flag(self): | |
| script = """ | |
| import torch | |
| """ | |
| try: | |
| subprocess.check_output( | |
| [sys.executable, '-OO', '-c', script], | |
| stderr=subprocess.STDOUT, | |
| # On Windows, opening the subprocess with the default CWD makes `import torch` | |
| # fail, so just set CWD to this script's directory | |
| cwd=os.path.dirname(os.path.realpath(__file__)),) | |
| except subprocess.CalledProcessError as e: | |
| self.assertFalse(e.returncode, "Import failed while running python in optimized mode") | |
| if __name__ == '__main__': | |
| run_tests() |