# Owner(s): ["module: nn"] import unittest from dataclasses import dataclass from functools import partial from itertools import chain, product import torch import torch.nn as nn import torch.nn.functional as F from torch.nn import CrossEntropyLoss from torch.nn.utils._expanded_weights import ExpandedWeight from torch.nn.utils._expanded_weights.expanded_weights_utils import ( forward_helper, set_grad_sample_if_exists, standard_kwargs, sum_over_all_but_batch_and_last_n, unpack_expanded_weight_or_tensor, ) from torch.nn.utils._per_sample_grad import call_for_per_sample_grads from torch.testing._internal.common_cuda import TEST_CUDA, tf32_off from torch.testing._internal.common_device_type import ( instantiate_device_type_tests, OpDTypes, ops, ) from torch.testing._internal.common_methods_invocations import op_db, SampleInput from torch.testing._internal.common_modules import module_db, modules from torch.testing._internal.common_nn import module_tests, new_module_tests, TestBase from torch.testing._internal.common_utils import ( freeze_rng_state, make_tensor, parametrize, run_tests, skipIfTorchDynamo, TestCase, ) from torch.utils._pytree import tree_map_only class TestContext: pass class TestExpandedWeightHelperFunction(TestCase): def test_forward_helper(self, device): input = torch.randn(3, 4, device=device) weight = torch.randn(5, 4, device=device) bias = torch.randn(5, device=device) for weight_batched, bias_batched in product([True, False], [True, False]): maybe_batched_weight = weight maybe_batched_bias = bias if weight_batched: maybe_batched_weight = ExpandedWeight( weight.clone().requires_grad_(), 3, loss_reduction="sum" ) if bias_batched: maybe_batched_bias = ExpandedWeight( bias.clone().requires_grad_(), 3, loss_reduction="sum" ) args = (input, maybe_batched_weight, maybe_batched_bias) expanded_args, expanded_kwargs = standard_kwargs(("bias",), args) res = forward_helper(nn.functional.linear, expanded_args, expanded_kwargs) expected = nn.functional.linear(input, weight, bias) self.assertEqual(res, expected) self.assertEqual(len(expanded_args), 2) assert expanded_args[0] is args[0] # avoids property checks in assertEquals assert expanded_args[1] is args[1] # avoids property checks in assertEquals self.assertEqual(len(expanded_kwargs), 1) assert ( expanded_kwargs["bias"] is args[2] ) # avoids property checks in assertEquals def test_forward_helper_failure_args(self, device): weight = torch.randn(5, 4, device=device) bias = torch.randn(5, device=device) with self.assertRaisesRegex( RuntimeError, r"do not support inputs that are also ExpandedWeights." ): input = ExpandedWeight( torch.randn(3, 4, requires_grad=True), 3, loss_reduction="sum" ) expanded_args, expanded_kwargs = standard_kwargs( ("bias",), (input, weight, bias) ) forward_helper(nn.functional.linear, expanded_args, expanded_kwargs) with self.assertRaisesRegex( RuntimeError, r"requires a Tensor as the first input" ): expanded_args, expanded_kwargs = standard_kwargs( ("bias",), (3, weight, bias) ) forward_helper(nn.functional.linear, expanded_args, expanded_kwargs) with self.assertRaisesRegex( RuntimeError, r"requires a batch dimension but got an input of size 0" ): expanded_args, expanded_kwargs = standard_kwargs( ("bias",), (torch.tensor(3), weight, bias) ) forward_helper(nn.functional.linear, expanded_args, expanded_kwargs) with self.assertRaisesRegex( RuntimeError, r"0 is not a valid batch size for Expanded Weights" ): expanded_args, expanded_kwargs = standard_kwargs( ("bias",), (torch.randn(0, 1, 2), weight, bias) ) forward_helper(nn.functional.linear, expanded_args, expanded_kwargs) input = torch.randn(3, 4) for weight_batched, bias_batched in product([True, False], [True, False]): if not weight_batched and not bias_batched: continue maybe_batched_weight = weight maybe_batched_bias = bias if weight_batched: maybe_batched_weight = ExpandedWeight( weight.clone().requires_grad_(), 4, loss_reduction="sum" ) if bias_batched: maybe_batched_bias = ExpandedWeight( bias.clone().requires_grad_(), 4, loss_reduction="sum" ) with self.assertRaisesRegex( RuntimeError, r"Expected ExpandedWeights to have batch size matching input", ): expanded_args, expanded_kwargs = standard_kwargs( ("bias",), (input, maybe_batched_weight, maybe_batched_bias) ) forward_helper(nn.functional.linear, expanded_args, expanded_kwargs) def test_set_grad_sample_if_exists(self, device): def test_fn(a): return grad_sample orig_weight = torch.randn(4, device=device, requires_grad=True) expanded_weight = ExpandedWeight(orig_weight, 3, loss_reduction="sum") grad_sample = torch.randn(3) set_grad_sample_if_exists(expanded_weight, test_fn) self.assertTrue(hasattr(orig_weight, "grad_sample")) self.assertEqual(orig_weight.grad_sample, grad_sample) basic_tensor = torch.randn(4, device=device) set_grad_sample_if_exists(basic_tensor, test_fn) self.assertFalse(hasattr(basic_tensor, "grad_sample")) non_tensor = 3 set_grad_sample_if_exists(non_tensor, test_fn) self.assertFalse(hasattr(non_tensor, "grad_sample")) def test_set_grad_sample_if_exists_failure(self, device): def test_fn(a): return True grad_tensor = torch.randn(4, requires_grad=True, device=device) with self.assertRaisesRegex( RuntimeError, r"does not support a mixture of ExpandedWeight parameters and normal Parameters", ): set_grad_sample_if_exists(grad_tensor, test_fn) def test_unpack_expanded_weight_or_tensor(self, device): input = torch.randn(3, requires_grad=True, device=device) self.assertEqual( input, unpack_expanded_weight_or_tensor( ExpandedWeight(input, 3, loss_reduction="sum") ), ) input.requires_grad_(False) self.assertEqual(input, unpack_expanded_weight_or_tensor(input)) self.assertTrue(unpack_expanded_weight_or_tensor(4) is None) def test_unpack_expanded_weight_or_tensor_with_custom_function(self, device): input = torch.randn(3, requires_grad=True, device=device) self.assertTrue( unpack_expanded_weight_or_tensor( ExpandedWeight(input, 3, loss_reduction="sum"), lambda x: x is input ) ) input.requires_grad_(False) self.assertTrue(unpack_expanded_weight_or_tensor(input, lambda x: x is input)) self.assertTrue( unpack_expanded_weight_or_tensor(4, lambda x: x is input) is None ) def test_unpack_expanded_weight_or_tensor_failure(self, device): input = torch.randn(3, requires_grad=True, device=device) with self.assertRaisesRegex( RuntimeError, r"does not support a mixture of ExpandedWeight parameters and normal Parameters", ): unpack_expanded_weight_or_tensor(input) with self.assertRaisesRegex( RuntimeError, r"does not support a mixture of ExpandedWeight parameters and normal Parameters", ): unpack_expanded_weight_or_tensor(input, lambda x: x is input) def test_sum_over_all_but_batch_and_last_n(self, device): input = torch.randn(1, 2, 3, 4, 5, device=device) res = sum_over_all_but_batch_and_last_n(input, 2) expected = input.sum((1, 2)) self.assertEqual(res, expected) res = sum_over_all_but_batch_and_last_n(input, 0) expected = input.sum((1, 2, 3, 4)) self.assertEqual(res, expected) res = sum_over_all_but_batch_and_last_n(input, 4) self.assertEqual(res, input) class TestExpandedWeightFunctional(TestCase): def _compare_ew_and_for_loop_per_sample_grads(self, op, sample_input, reduction): input = sample_input.input args = sample_input.args kwargs = sample_input.kwargs batch_size = input.shape[0] if len(input.shape) > 1 else 1 # get per sample grads with ExpandedWeights objects loss_reduction = "sum" if reduction == torch.sum else "mean" (ew_input, ew_args, ew_kwargs) = make_expanded_weight( sample_input, batch_size, loss_reduction ) diff_input_list = (ew_input,) + tuple(ew_args) + tuple(ew_kwargs.values()) diff_input_list = [i for i in diff_input_list if is_diff_tensor(i)] diff_input_list = [ i.orig_weight if isinstance(i, ExpandedWeight) else i for i in diff_input_list ] if not diff_input_list: return result = run_op(op, ew_input, *ew_args, **ew_kwargs) reduction( result ).backward() # grad doesn't work with ExpandedWeight because it calls __torch_function__ expanded_weight_grad = tuple( i.grad_sample if hasattr(i, "grad_sample") else i.grad for i in diff_input_list ) # get per sample grads with for loop func = partial(run_op, op) per_sample_grad = for_loop_per_sample_grad( batch_size, reduction, input, func, *args, **kwargs ) # check equality self.assertEqual(len(per_sample_grad), len(expanded_weight_grad)) if loss_reduction == "mean": # don't check equality of `input.grad`s since these vanilla tensors won't be scaled expanded_weight_grad = expanded_weight_grad[1:] per_sample_grad = per_sample_grad[1:] for result_grad, expected_grad in zip(expanded_weight_grad, per_sample_grad): self.assertEqual(result_grad, expected_grad) @ops( filter(lambda op: op.supports_expanded_weight, op_db), dtypes=OpDTypes.supported, allowed_dtypes=(torch.double,), ) def test_expanded_weight_per_sample_grad_sum(self, device, dtype, op): sample_inputs = op.sample_inputs(device, dtype, requires_grad=True) for sample_input in supported_inputs(op, sample_inputs): if ( op.name == "nn.functional.embedding" ): # embedding flips its argument order for autograd tests sample_input = SampleInput( sample_input.args[0], args=(sample_input.input,), kwargs=sample_input.kwargs, ) self._compare_ew_and_for_loop_per_sample_grads(op, sample_input, torch.sum) @ops( filter(lambda op: op.supports_expanded_weight, op_db), dtypes=OpDTypes.supported, allowed_dtypes=(torch.double,), ) def test_expanded_weight_per_sample_grad_mean(self, device, dtype, op): sample_inputs = op.sample_inputs(device, dtype, requires_grad=True) for sample_input in supported_inputs(op, sample_inputs): if ( op.name == "nn.functional.embedding" ): # embedding flips its argument order for autograd tests sample_input = SampleInput( sample_input.args[0], args=(sample_input.input,), kwargs=sample_input.kwargs, ) self._compare_ew_and_for_loop_per_sample_grads(op, sample_input, torch.mean) @ops( filter(lambda op: op.supports_expanded_weight, op_db), dtypes=OpDTypes.supported, allowed_dtypes=(torch.double,), ) def test_expanded_weights_per_sample_grad_input_no_grad(self, device, dtype, op): sample_inputs = op.sample_inputs(device, dtype, requires_grad=True) for sample_input in supported_inputs(op, sample_inputs): if ( op.name == "nn.functional.embedding" ): # embedding flips its argument order for autograd tests sample_input = SampleInput( sample_input.args[0], args=(sample_input.input,), kwargs=sample_input.kwargs, ) sample_input.input.requires_grad_(False) self._compare_ew_and_for_loop_per_sample_grads(op, sample_input, torch.mean) @skipIfTorchDynamo("Checking error message doesn't work with dynamo") @ops( filter(lambda op: op.supports_expanded_weight, op_db), dtypes=OpDTypes.supported, allowed_dtypes=(torch.double,), ) def test_unsupported_expand_weights(self, device, dtype, op): sample_inputs = op.sample_inputs(device, dtype, requires_grad=True) unsupported_inputs = supported_inputs(op, sample_inputs, supported_inputs=False) for sample_input in unsupported_inputs: with self.assertRaisesRegex(RuntimeError, r"Expanded Weights"): if ( op.name == "nn.functional.embedding" ): # embedding flips its argument order for autograd tests sample_input = SampleInput( sample_input.args[0], args=(sample_input.input,), kwargs=sample_input.kwargs, ) input = sample_input.input batch_size = input.shape[0] if len(input.shape) > 1 else 1 # get per sample grads with ExpandedWeights objects (ew_input, ew_args, ew_kwargs) = make_expanded_weight( sample_input, batch_size ) result = run_op(op, ew_input, *ew_args, **ew_kwargs) diff_input_list = ( (ew_input,) + tuple(ew_args) + tuple(ew_kwargs.values()) ) diff_input_list = [i for i in diff_input_list if is_diff_tensor(i)] diff_input_list = [ i.orig_weight if isinstance(i, ExpandedWeight) else i for i in diff_input_list ] result.sum().backward() # grad doesn't work with ExpandedWeight because it calls __torch_function__ @ops( filter(lambda op: op.supports_expanded_weight, op_db), dtypes=OpDTypes.supported ) def test_expanded_weight_forward(self, device, dtype, op): sample_inputs = op.sample_inputs(device, dtype) for sample_input in supported_inputs(op, sample_inputs): if ( op.name == "nn.functional.embedding" ): # embedding flips its argument order for autograd tests sample_input = SampleInput( sample_input.args[0].clone(), args=(sample_input.input.clone(),), kwargs=sample_input.kwargs, ) if ( "cuda" in device and "max_norm" in sample_input.kwargs and "padding_idx" in sample_input.kwargs ): self.skipTest( "embedding is non-determinstic in this case, see issue #74679" ) batch_size = ( sample_input.input.shape[0] if len(sample_input.input.shape) > 1 else 1 ) for loss_reduction in ["sum", "mean"]: (ew_input, ew_args, ew_kwargs) = make_expanded_weight( sample_input, batch_size, loss_reduction ) expanded_weight_result = run_op(op, ew_input, *ew_args, **ew_kwargs) normal_result = run_op( op, sample_input.input, *sample_input.args, **sample_input.kwargs ) self.assertEqual(expanded_weight_result, normal_result) def test_expanded_weight_error(self, device): batch_size = 3 sample_input = make_tensor( (batch_size, 4), dtype=torch.float32, device=device, requires_grad=True ) sample_weight = make_tensor( (4), dtype=torch.float32, device=device, requires_grad=True ) with self.assertRaisesRegex( RuntimeError, r"Expanded Weights encountered but cannot handle function" ): torch.add( sample_input, ExpandedWeight(sample_weight, batch_size, loss_reduction="sum"), ) def _test_embedding_model(self, model, num_embedding, device): batch_size = 32 input = torch.randint(0, num_embedding, (batch_size, 5, 5), device=device) return self._test_model( partial(model, num_embedding=num_embedding), batch_size, input, device ) def _test_conv_model( self, model, input_size, num_dim, device, loss_reduction="sum", atol=1e-4, rtol=5e-5, ): batch_size = 32 input_ending = [input_size] * num_dim input = torch.randn([batch_size, 3] + input_ending, device=device) return self._test_model( partial(model, num_dim=num_dim), batch_size, input, device, loss_reduction, atol, rtol, ) def _test_model( self, model, batch_size, input, device, loss_reduction="sum", atol=1e-4, rtol=5e-5, ): model = model(10).to(device) targets = torch.randint(0, 10, (batch_size,), device=device) criterion = CrossEntropyLoss(reduction=loss_reduction) result = call_for_per_sample_grads(model, loss_reduction=loss_reduction)(input) loss = criterion(result, targets) loss.backward() result = [] for weight in model.parameters(): result.append(weight.grad_sample) del weight.grad_sample expected = [] for i in range(batch_size): loss = criterion(model(input[i].unsqueeze(0)), targets[i].unsqueeze(0)) expected.append( torch.autograd.grad(loss, model.parameters(), torch.ones_like(loss)) ) expected = [torch.stack(grad) for grad in zip(*expected)] for res, exp in zip(result, expected): self.assertEqual(res, exp, atol=atol, rtol=rtol) def _compute_tolerances(self, device): is_cuda_sm86 = device.startswith("cuda") and torch.cuda.get_device_capability( 0 ) == (8, 6) return (9e-3, 5e-5) if is_cuda_sm86 else (1e-4, 5e-5) @tf32_off() def test_cnn_model_sum(self, device): def convnet(num_classes, num_dim): return nn.Sequential( nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.AvgPool2d(kernel_size=2, stride=2), nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.AvgPool2d(kernel_size=2, stride=2), nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.AvgPool2d(kernel_size=2, stride=2), nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(start_dim=1, end_dim=-1), nn.Linear(128, num_classes, bias=True), ) atol, rtol = self._compute_tolerances(device) return self._test_conv_model(convnet, 28, 2, device, atol=atol, rtol=rtol) @tf32_off() def test_cnn_model_mean(self, device): def convnet(num_classes, num_dim): return nn.Sequential( nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.AvgPool2d(kernel_size=2, stride=2), nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.AvgPool2d(kernel_size=2, stride=2), nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.AvgPool2d(kernel_size=2, stride=2), nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(start_dim=1, end_dim=-1), nn.Linear(128, num_classes, bias=True), ) atol, rtol = self._compute_tolerances(device) return self._test_conv_model( convnet, 28, 2, device, loss_reduction="mean", atol=atol, rtol=rtol ) @parametrize("num_dim", [1, 2, 3]) @tf32_off() def test_instance_norm_model(self, num_dim, device): def instance_norm_model(num_classes, num_dim): conv_layer = ( nn.Conv1d if num_dim == 1 else nn.Conv2d if num_dim == 2 else nn.Conv3d ) norm_layer = ( nn.InstanceNorm1d if num_dim == 1 else nn.InstanceNorm2d if num_dim == 2 else nn.InstanceNorm3d ) return nn.Sequential( conv_layer(3, 32, kernel_size=3, stride=1, padding=1), norm_layer(32, affine=True), nn.Flatten(start_dim=1, end_dim=-1), nn.Linear(32 * (7**num_dim), num_classes, bias=True), ) atol, rtol = self._compute_tolerances(device) return self._test_conv_model( instance_norm_model, 7, num_dim, device, atol=atol, rtol=rtol ) @parametrize("num_dim", [1, 2, 3]) @tf32_off() def test_group_norm_model(self, num_dim, device): def group_norm_model(num_classes, num_dim): conv_layer = ( nn.Conv1d if num_dim == 1 else nn.Conv2d if num_dim == 2 else nn.Conv3d ) return nn.Sequential( conv_layer(3, 32, kernel_size=3, stride=1, padding=1), nn.GroupNorm(8, 32, affine=True), nn.Flatten(start_dim=1, end_dim=-1), nn.Linear(32 * (7**num_dim), num_classes, bias=True), ) atol, rtol = self._compute_tolerances(device) return self._test_conv_model( group_norm_model, 7, num_dim, device, atol=atol, rtol=rtol ) @parametrize("num_dim", [1, 2, 3]) @tf32_off() def test_layer_norm_model(self, num_dim, device): def layer_norm_model(num_classes, num_dim): conv_layer = ( nn.Conv1d if num_dim == 1 else nn.Conv2d if num_dim == 2 else nn.Conv3d ) normalized_shape = [7] * num_dim return nn.Sequential( conv_layer(3, 32, kernel_size=3, stride=1, padding=1), nn.LayerNorm(normalized_shape, elementwise_affine=True), nn.Flatten(start_dim=1, end_dim=-1), nn.Linear(32 * (7**num_dim), num_classes, bias=True), ) atol, rtol = self._compute_tolerances(device) return self._test_conv_model( layer_norm_model, 7, num_dim, device, atol=atol, rtol=rtol ) def test_embedding_model(self, device): def embedding_model(num_classes, num_embedding): return nn.Sequential( nn.Embedding(num_embedding, 15), nn.Flatten(start_dim=1, end_dim=-1), nn.Linear(375, num_classes, bias=True), ) return self._test_embedding_model(embedding_model, 16, device) def test_group_norm_error(self, device): # group norm has to call native_group_norm. This checks that it hits the same errors # that normal group norm would N = 3 C = 5 inp = torch.randn(N, C) with self.assertRaisesRegex( RuntimeError, r"Expected number of channels in input to be divisible" ): F.group_norm(inp, 2) # 5 is not divisible by 2 class TestExpandedWeightModule(TestCase): def _do_test( self, module, input, args=None, kwargs=None, batch_first=True, atol=None, rtol=None, ): args = args or () kwargs = kwargs or {} batch_dim = 0 if batch_first else 1 batch_size = input.shape[batch_dim] diff_input = input.dtype == torch.float or input.dtype == torch.double if diff_input: input.requires_grad_() with freeze_rng_state(): # get per sample grads with ExpandedWeights context manager actual_res = call_for_per_sample_grads( module, batch_size=batch_size, loss_reduction="sum", batch_first=batch_first, )(input, *args, **kwargs).sum() actual_res.backward() actual_grads = [] for param in module.parameters(): actual_grads.append(param.grad_sample) del param.grad_sample if diff_input: actual_grads.append(input.grad.clone()) input.grad = torch.zeros_like(input.grad) # get per sample grads with a for loop expected_res = torch.tensor( 0.0, device=input.device, dtype=actual_res.dtype ) expected_grads = [] for i in range(batch_size): input_slice = input.narrow(batch_dim, i, 1) input_slice = input_slice.squeeze(batch_dim) # h's batch dim is always the first dim. Must be contiguous for CUDA sliced_args = tree_map_only( torch.Tensor, lambda t: t.narrow(1, i, 1).contiguous(), args ) diff_params = module.parameters() if diff_input: diff_params = chain(diff_params, (input_slice,)) res = module( input_slice.unsqueeze(batch_dim).contiguous(), *sliced_args, **kwargs, ).sum() out_grads = torch.autograd.grad( res, diff_params, torch.ones_like(res), allow_unused=True ) expected_grads.append(out_grads) expected_res += res expected_grads = [torch.stack(grad) for grad in zip(*expected_grads)] if not batch_first: expected_grads[-1] = expected_grads[-1].transpose(0, 1) self.assertEqual(actual_res, expected_res) [ self.assertEqual(actual, expected, atol=atol, rtol=rtol) for (actual, expected) in zip(actual_grads, expected_grads) ] def _do_test_multi_input(self, module, input): class TestModule(nn.Module): def __init__(self, module): super().__init__() self.module = module def forward(self, input): return self.module(input) + self.module(input) batch_size = input.shape[0] diff_input = input.dtype == torch.float or input.dtype == torch.double if diff_input: input.requires_grad_() with freeze_rng_state(): # get per sample grads with ExpandedWeights context manager, calling .backward() twice test_module = TestModule(module) actual_res = call_for_per_sample_grads(test_module, loss_reduction="sum")( input ).sum() actual_res.backward() actual_grads = [] for param in module.parameters(): actual_grads.append(param.grad_sample) del param.grad_sample if diff_input: actual_grads.append(input.grad.clone()) input.grad = torch.zeros_like(input.grad) # get per sample grads with a for loop, running over the input twice expected_grads = [] for i in range(batch_size): input_slice = input[i] diff_params = module.parameters() if diff_input: diff_params = chain(diff_params, (input_slice,)) res = module(input_slice.unsqueeze(0)).sum() out_grads = torch.autograd.grad( res, diff_params, torch.ones_like(res), allow_unused=True ) expected_grads.append(out_grads) expected_grads = tuple(torch.stack(grad) for grad in zip(*expected_grads)) expected_grads = tuple( expected_grad for expected_grad in expected_grads if expected_grad is not None ) assert [ self.assertEqual(actual, 2 * expected) for (actual, expected) in zip(actual_grads, expected_grads) ] def _do_test_rnn_packed_sequence( self, module, input, args=None, kwargs=None, atol=None, rtol=None ): args = args if args is not None else () kwargs = kwargs if kwargs is not None else {} batch_size = max(tuple(input.batch_sizes)).item() with freeze_rng_state(): # get per sample grads with ExpandedWeights context manager actual_res = call_for_per_sample_grads( module, batch_size=batch_size, loss_reduction="sum" )(input, *args, **kwargs).data.sum() actual_res.backward() actual_grads = [] for param in module.parameters(): self.assertEqual(param.grad_sample.shape[0], batch_size) actual_grads.append(param.grad_sample) del param.grad_sample input.data.grad = torch.zeros_like(input.data) # compute the per sample grads with a for loop expected_res = torch.zeros_like(actual_res) expected_grads = [] padded_input, seq_sizes = torch.nn.utils.rnn.pad_packed_sequence( input, batch_first=True ) for i in range(len(seq_sizes)): input_slice = padded_input[i].narrow(0, 0, seq_sizes[i]) diff_params = module.parameters() batch_dim = 0 if module.m.batch_first else 1 res = module(input_slice.unsqueeze(batch_dim), *args, **kwargs).sum() expected_res += res out_grads = torch.autograd.grad( res, diff_params, torch.ones_like(res), allow_unused=True ) expected_grads.append(out_grads) expected_grads = [torch.stack(grad) for grad in zip(*expected_grads)] self.assertEqual(actual_res, expected_res) [ self.assertEqual(actual, expected, atol=atol, rtol=rtol) for (actual, expected) in zip(actual_grads, expected_grads) ] @modules( filter( lambda m_info: m_info.module_cls in (torch.nn.RNN, torch.nn.LSTM, torch.nn.GRU), module_db, ) ) @tf32_off() def test_module(self, device, dtype, module_info, training): class RNNWrapper(torch.nn.Module): def __init__(self, m_cons, args, kwargs): super().__init__() self.m = m_cons(*args, **kwargs) def forward(self, *inps): ret = self.m(*inps) assert isinstance(ret, tuple) return ret[0] def batch_hidden(h): new_h_shape = [1] * (len(h.shape) + 1) new_h_shape[1] = 2 return h.unsqueeze(1).repeat(new_h_shape) module_cls = module_info.module_cls atol, rtol = ( (1e-4, 1e-5) if module_cls == torch.nn.GRU and dtype == torch.float32 else (None, None) ) module_inputs = module_info.module_inputs_func( module_info, device=device, dtype=dtype, requires_grad=True, training=training, with_packed_sequence=True, ) for module_input in module_inputs: if module_input.forward_input is None: continue args, kwargs = ( module_input.constructor_input.args, module_input.constructor_input.kwargs, ) m = RNNWrapper(module_cls, args, kwargs) batch_first = m.m.batch_first m.to(device).to(dtype) args, kwargs = ( module_input.forward_input.args, module_input.forward_input.kwargs, ) # if the RNN tests use unbatched inputs--batch the inputs input = args[0] if isinstance(input, torch.Tensor) and input.dim() == 2: input = input.detach() new_input_shape = [1] * (len(input.shape) + 1) if batch_first: new_input_shape[0] = 2 input = input.repeat(new_input_shape) else: new_input_shape[1] = 2 input = input.unsqueeze(1).repeat(new_input_shape) h = args[1] if len(args) > 1 else None if h is not None: h = ( batch_hidden(h) if isinstance(h, torch.Tensor) else tuple(batch_hidden(hx) for hx in h) ) args = list(args) args[1] = h if isinstance(input, torch.nn.utils.rnn.PackedSequence): self._do_test_rnn_packed_sequence( m, input, args[1:], kwargs, atol=atol, rtol=rtol ) else: self._do_test( m, input, args[1:], kwargs, batch_first=batch_first, atol=atol, rtol=rtol, ) def test_per_sample_api_failing(self): module = nn.Linear(10, 10) input = torch.randn(64, 10) with self.assertRaisesRegex(RuntimeError, r"Module passed must be nn.Module"): call_for_per_sample_grads("fail")(input) with self.assertRaisesRegex( RuntimeError, r"Batch size passed must be None or an integer" ): call_for_per_sample_grads(module, batch_size=6.4)(input) with self.assertRaisesRegex(RuntimeError, r"Batch size must be positive"): call_for_per_sample_grads(module, batch_size=-64)(input) with self.assertRaisesRegex(RuntimeError, r"incorrect for multiple calls"): loss = call_for_per_sample_grads(module)(input).sum() loss.backward() # populate grad_sample fields call_for_per_sample_grads(module)(input) module = nn.Linear(10, 10) # reset to not have grad_sample fields with self.assertRaisesRegex( RuntimeError, r"Expected loss_reduction argument to be sum or mean" ): call_for_per_sample_grads(module, loss_reduction="")(input) def test_per_sample_api_compute_batch_size(self): class CustomModule(nn.Module): def __init__(self) -> None: super().__init__() self.linear = nn.Linear(5, 5) def forward(self, input1, input2): return self.linear(input1) + self.linear(input2) module = CustomModule() input1 = torch.randn(4, 5) input2 = torch.randn(5, 5) with self.assertRaisesRegex( RuntimeError, "found at least one input with batch size 4 and one with batch size 5", ): call_for_per_sample_grads(module)(input1, input2) input2 = torch.randn(4, 5) call_for_per_sample_grads(module)(input1, input2) module = CustomModule() call_for_per_sample_grads(module)(input1, input2=input2) module = CustomModule() call_for_per_sample_grads(module)(input1=input1, input2=input2) def test_per_sample_api_compute_batch_size_not_pytreeable(self): @dataclass class NonPytreeableTuple: elem1: torch.Tensor elem2: torch.Tensor class CustomModule(nn.Module): def __init__(self) -> None: super().__init__() self.linear = nn.Linear(5, 5) def forward(self, input1, input2): return self.linear(input1.elem1) + self.linear(input1.elem2) input = NonPytreeableTuple(torch.randn(4, 5), torch.randn(4, 5)) model = CustomModule() with self.assertRaisesRegex( RuntimeError, "ExpandedWeights cannot compute the batch size from the inputs", ): call_for_per_sample_grads(model)(input, "") # would prefer for it to error because input is not pytree-able but that's hard to detect with self.assertRaisesRegex( RuntimeError, "Expected ExpandedWeights to have batch size matching input" ): call_for_per_sample_grads(model)(input, torch.randn(5)) model = CustomModule() # TODO: functional call bug, sam will fix call_for_per_sample_grads(model)(input, torch.randn(4, 5)) model = CustomModule() call_for_per_sample_grads(model, batch_size=4)(input, torch.randn(5)) class ContextManagerTests(TestBase): def __init__(self, *args, **kwargs): self.test_cpu = kwargs.get("test_cpu", True) self.test_cuda = kwargs.get("test_cuda", True) super().__init__(*args, **kwargs) @property def constructor_args(self): return self._get_arg("constructor_args", False) def test_context_manager(self, test_case, device): kwargs = {"device": device, "dtype": torch.double} module = self.constructor(*self.constructor_args).to(**kwargs) if "Embedding" in self.get_name(): kwargs["dtype"] = torch.long input = self._get_input().to(**kwargs) if len(input.shape) == 0 or input.shape[0] == 0: raise unittest.SkipTest( "Can't get per sample gradients when no batch dim or batch dim is 0" ) if self.constructor == torch.nn.Linear and len(input.shape) == 1: raise unittest.SkipTest( "Can't get per sample gradients for input of rank 1" ) test_case._do_test(module, input) def test_context_manager_multiple_inputs(self, test_case, device): module = self.constructor(*self.constructor_args).to(device) input = self._get_input() if len(input.shape) == 0 or input.shape[0] == 0: raise unittest.SkipTest( "Can't get per sample gradients when no batch dim or batch dim is 0" ) if self.constructor == torch.nn.Linear and len(input.shape) == 1: raise unittest.SkipTest( "Can't get per sample gradients for input of rank 1" ) test_case._do_test_multi_input(module, input) def filter_supported_tests(t): supported_modules = [ "Linear", "Conv1d", "Conv2d", "Conv3d", "Embedding", "LayerNorm", "GroupNorm", "InstanceNorm", ] if "module_name" in t and t["module_name"] in supported_modules: return True # TODO: Once all of these use ModuleInfo, replace with ModuleInfo tests # These currently use the legacy nn tests supported_tests = [ t for t in module_tests + new_module_tests if filter_supported_tests(t) ] for test_param in supported_tests: if "constructor" not in test_param: name = test_param.pop("module_name") test_param["constructor"] = getattr(nn, name) decorator = test_param.pop("decorator", lambda test: test) test = ContextManagerTests(**test_param) test_name = test.get_name() if hasattr(TestExpandedWeightModule, test_name): raise RuntimeError("Found two tests with the same name: " + test_name) test_name_multi_input = test.get_name() + "_multiple_inputs" if hasattr(TestExpandedWeightModule, test_name_multi_input): raise RuntimeError("Found two tests with the same name: " + test_name) if test.test_cpu: setattr( TestExpandedWeightModule, test_name, decorator(lambda self, test=test: test.test_context_manager(self, "cpu")), ) setattr( TestExpandedWeightModule, test_name_multi_input, decorator( lambda self, test=test: test.test_context_manager_multiple_inputs( self, "cpu" ) ), ) if TEST_CUDA and test.test_cuda: # since this checks derivatives, only use double for precision setattr( TestExpandedWeightModule, test_name + "_cuda_double", decorator(lambda self, test=test: test.test_context_manager(self, "cuda")), ) # ------------- HELPER FUNCTIONS ----------------- def run_op(op, input, *args, **kwargs): r""" OpInfo for Embedding switches the input and weight so autograd tests will only check the derivative of the weight, not the input, which can't be differentiable since its dtype is int. Calls op, using the special ordering that Embedding's OpInfo expects for that case. """ if op.name == "nn.functional.embedding": return op(args[0], input, **kwargs) else: return op(input, *args, **kwargs) def make_expanded_weight(sample_input, batch_size, loss_reduction="sum"): def expanded_weight_or_clone(arg): if is_diff_tensor(arg): return ExpandedWeight(torch.clone(arg), batch_size, loss_reduction) return clone_if_tensor(arg) ew_input = clone_if_tensor(sample_input.input) ew_args = tuple(expanded_weight_or_clone(arg) for arg in sample_input.args) ew_kwargs = { name: expanded_weight_or_clone(arg) for (name, arg) in sample_input.kwargs.items() } return ew_input, ew_args, ew_kwargs def supported_inputs(op, sample_inputs, supported_inputs=True): r""" ExpandedWeights currently does not support some use cases when there's no batch dimension or operations that would cause inter-batch operations. Removes all of the cases it cannot deal with """ def filter_fn(input): convolutions = [ "nn.functional.conv1d", "nn.functional.conv2d", "nn.functional.conv3d", ] batched_input_size = dict(zip(convolutions, [3, 4, 5])) if op.name == "nn.functional.linear": is_supported_input = ( input.input.dim() > 1 ) # input of rank 1 means no batch dim elif op.name == "nn.functional.layer_norm": normalized_shape = input.args[0] is_supported_input = ( input.input.shape != normalized_shape ) # would cause inter-batch operations elif op.name in convolutions: # currently can't deal with padding computation on Python level is_supported_input = input.input.dim() == batched_input_size[op.name] elif op.name == "nn.functional.embedding": idx = input.args[0] is_supported_input = len(idx.shape) > 1 # there's no batch size else: is_supported_input = True is_supported_input = ( is_supported_input and input.input.shape[0] > 0 ) # 0 is not a valid batch size return is_supported_input if supported_inputs else not is_supported_input return [input for input in sample_inputs if filter_fn(input)] def for_loop_per_sample_grad(batch_size, reduction, input, func, *args, **kwargs): # get per sample grads by getting derivative for each input in a for loop per_sample_grad = [] for i in range(batch_size): per_sample_input = input[i] result = reduction(func(per_sample_input.unsqueeze(0), *args, **kwargs)) diff_input_list = (per_sample_input,) + tuple(args) + tuple(kwargs.values()) diff_input_list = [ i for i in diff_input_list if isinstance(i, torch.Tensor) and i.requires_grad ] per_sample_grad.append( torch.autograd.grad( result, diff_input_list, torch.ones_like(result), allow_unused=True ) ) if len(per_sample_grad) == batch_size: per_sample_grad = tuple(torch.stack(grad) for grad in zip(*per_sample_grad)) return per_sample_grad def is_diff_tensor(t): return isinstance(t, ExpandedWeight) or ( isinstance(t, torch.Tensor) and t.requires_grad ) def clone_if_tensor(t): if isinstance(t, torch.Tensor): res = torch.clone(t).detach() res.requires_grad_(t.requires_grad) return res else: return t instantiate_device_type_tests(TestExpandedWeightHelperFunction, globals()) instantiate_device_type_tests(TestExpandedWeightFunctional, globals()) instantiate_device_type_tests(TestExpandedWeightModule, globals()) if __name__ == "__main__": run_tests()