Merge branch 'main' into linear-gramian-computer

Author: ValerianRey

tests/speed/autogram/grad_vs_jac_vs_gram.py

@@ -14,7 +14,6 @@
     NoFreeParam,
     SqueezeNet,
     WithTransformerLarge,
-    get_in_out_shapes,
 )
 from utils.forward_backwards import (
     autograd_forward_backward,
@@ -23,7 +22,7 @@
     autojac_forward_backward,
     make_mse_loss_fn,
 )
-from utils.tensors import make_tensors
+from utils.tensors import make_inputs_and_targets
 
 from torchjd.aggregation import Mean
 from torchjd.autogram import Engine
@@ -43,9 +42,7 @@
 
 def compare_autograd_autojac_and_autogram_speed(factory: ModuleFactory, batch_size: int):
     model = factory()
-    input_shapes, output_shapes = get_in_out_shapes(model)
-    inputs = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model, batch_size)
     loss_fn = make_mse_loss_fn(targets)
 
     A = Mean()
@@ -64,7 +61,7 @@ def init_fn_autograd():
         fn_autograd()
 
     def fn_autograd_gramian():
-        autograd_gramian_forward_backward(model, inputs, list(model.parameters()), loss_fn, W)
+        autograd_gramian_forward_backward(model, inputs, loss_fn, W)
 
     def init_fn_autograd_gramian():
         torch.cuda.empty_cache()
@@ -80,7 +77,7 @@ def init_fn_autojac():
         fn_autojac()
 
     def fn_autogram():
-        autogram_forward_backward(model, engine, W, inputs, loss_fn)
+        autogram_forward_backward(model, inputs, loss_fn, engine, W)
 
     def init_fn_autogram():
         torch.cuda.empty_cache()

tests/unit/autogram/test_engine.py

@@ -6,7 +6,7 @@
 import torch
 from pytest import mark, param
 from torch import Tensor
-from torch.nn import RNN, BatchNorm2d, InstanceNorm2d, Linear
+from torch.nn import BatchNorm2d, InstanceNorm2d, Linear
 from torch.optim import SGD
 from torch.testing import assert_close
 from utils.architectures import (
@@ -56,25 +56,26 @@
     WithModuleWithStringOutput,
     WithMultiHeadAttention,
     WithNoTensorOutput,
+    WithRNN,
     WithSideEffect,
     WithSomeFrozenModule,
     WithTransformer,
     WithTransformerLarge,
-    get_in_out_shapes,
 )
 from utils.dict_assertions import assert_tensor_dicts_are_close
 from utils.forward_backwards import (
     autograd_forward_backward,
     autogram_forward_backward,
     compute_gramian,
     compute_gramian_with_autograd,
+    forward_pass,
     make_mse_loss_fn,
     reduce_to_first_tensor,
     reduce_to_matrix,
     reduce_to_scalar,
     reduce_to_vector,
 )
-from utils.tensors import make_tensors, ones_, randn_, zeros_
+from utils.tensors import make_inputs_and_targets, ones_, randn_, zeros_
 
 from torchjd.aggregation import UPGradWeighting
 from torchjd.autogram._engine import Engine
@@ -143,22 +144,14 @@ def _assert_gramian_is_equivalent_to_autograd(
     factory: ModuleFactory, batch_size: int, batch_dim: int | None
 ):
     model_autograd, model_autogram = factory(), factory()
-    input_shapes, output_shapes = get_in_out_shapes(model_autograd)
-
     engine = Engine(model_autogram, batch_dim=batch_dim)
-
-    inputs = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model_autograd, batch_size)
     loss_fn = make_mse_loss_fn(targets)
 
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_autograd(inputs)
-    losses = reduce_to_vector(loss_fn(output))
+    losses = forward_pass(model_autograd, inputs, loss_fn, reduce_to_vector)
     autograd_gramian = compute_gramian_with_autograd(losses, list(model_autograd.parameters()))
 
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_autogram(inputs)
-    losses = reduce_to_vector(loss_fn(output))
+    losses = forward_pass(model_autogram, inputs, loss_fn, reduce_to_vector)
     autogram_gramian = engine.compute_gramian(losses)
 
     assert_close(autogram_gramian, autograd_gramian, rtol=1e-4, atol=3e-5)
@@ -179,10 +172,7 @@ def test_compute_gramian(factory: ModuleFactory, batch_size: int, batch_dim: int
         ModuleFactory(WithSideEffect),
         ModuleFactory(Randomness),
         ModuleFactory(InstanceNorm2d, num_features=3, affine=True, track_running_stats=True),
-        param(
-            ModuleFactory(RNN, input_size=8, hidden_size=5, batch_first=True),
-            marks=mark.xfail_if_cuda,
-        ),
+        param(ModuleFactory(WithRNN), marks=mark.xfail_if_cuda),
     ],
 )
 @mark.parametrize("batch_size", [1, 3, 32])
@@ -257,26 +247,18 @@ def test_compute_gramian_various_output_shapes(
 
     factory = ModuleFactory(Ndim2Output)
     model_autograd, model_autogram = factory(), factory()
-    input_shapes, output_shapes = get_in_out_shapes(model_autograd)
-
-    engine = Engine(model_autogram, batch_dim=batch_dim)
-
-    inputs = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model_autograd, batch_size)
     loss_fn = make_mse_loss_fn(targets)
 
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_autograd(inputs)
-    losses = reduction(loss_fn(output))
+    losses = forward_pass(model_autograd, inputs, loss_fn, reduction)
     reshaped_losses = torch.movedim(losses, movedim_source, movedim_destination)
     # Go back to a vector so that compute_gramian_with_autograd works
     loss_vector = reshaped_losses.reshape([-1])
     autograd_gramian = compute_gramian_with_autograd(loss_vector, list(model_autograd.parameters()))
     expected_gramian = reshape_gramian(autograd_gramian, list(reshaped_losses.shape))
 
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_autogram(inputs)
-    losses = reduction(loss_fn(output))
+    engine = Engine(model_autogram, batch_dim=batch_dim)
+    losses = forward_pass(model_autogram, inputs, loss_fn, reduction)
     reshaped_losses = torch.movedim(losses, movedim_source, movedim_destination)
     autogram_gramian = engine.compute_gramian(reshaped_losses)
 
@@ -298,30 +280,20 @@ def test_compute_partial_gramian(gramian_module_names: set[str], batch_dim: int
     the model parameters is specified.
     """
 
-    factory = ModuleFactory(SimpleBranched)
-    model = factory()
-    input_shapes, output_shapes = get_in_out_shapes(model)
+    model = SimpleBranched()
     batch_size = 64
-
-    input = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model, batch_size)
     loss_fn = make_mse_loss_fn(targets)
-
-    output = model(input)
-    losses = reduce_to_vector(loss_fn(output))
-
     gramian_modules = [model.get_submodule(name) for name in gramian_module_names]
     gramian_params = []
     for m in gramian_modules:
         gramian_params += list(m.parameters())
 
+    losses = forward_pass(model, inputs, loss_fn, reduce_to_vector)
     autograd_gramian = compute_gramian_with_autograd(losses, gramian_params, retain_graph=True)
-    torch.manual_seed(0)
 
     engine = Engine(*gramian_modules, batch_dim=batch_dim)
-
-    output = model(input)
-    losses = reduce_to_vector(loss_fn(output))
+    losses = forward_pass(model, inputs, loss_fn, reduce_to_vector)
     gramian = engine.compute_gramian(losses)
 
     assert_close(gramian, autograd_gramian)
@@ -333,22 +305,15 @@ def test_iwrm_steps_with_autogram(factory: ModuleFactory, batch_size: int, batch
     """Tests that the autogram engine doesn't raise any error during several IWRM iterations."""
 
     n_iter = 3
-
     model = factory()
-    input_shapes, output_shapes = get_in_out_shapes(model)
-
     weighting = UPGradWeighting()
-
     engine = Engine(model, batch_dim=batch_dim)
     optimizer = SGD(model.parameters(), lr=1e-7)
 
     for i in range(n_iter):
-        inputs = make_tensors(batch_size, input_shapes)
-        targets = make_tensors(batch_size, output_shapes)
+        inputs, targets = make_inputs_and_targets(model, batch_size)
         loss_fn = make_mse_loss_fn(targets)
-
-        autogram_forward_backward(model, engine, weighting, inputs, loss_fn)
-
+        autogram_forward_backward(model, inputs, loss_fn, engine, weighting)
         optimizer.step()
         model.zero_grad()
 
@@ -365,29 +330,22 @@ def test_autograd_while_modules_are_hooked(
     """
 
     model, model_autogram = factory(), factory()
-    input_shapes, output_shapes = get_in_out_shapes(model)
-
-    input = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model, batch_size)
     loss_fn = make_mse_loss_fn(targets)
 
-    torch.manual_seed(0)  # Fix randomness for random models
-    autograd_forward_backward(model, input, loss_fn)
+    autograd_forward_backward(model, inputs, loss_fn)
     autograd_grads = {name: p.grad for name, p in model.named_parameters() if p.grad is not None}
 
     # Hook modules and optionally compute the Gramian
     engine = Engine(model_autogram, batch_dim=batch_dim)
     if use_engine:
-        torch.manual_seed(0)  # Fix randomness for random models
-        output = model_autogram(input)
-        losses = reduce_to_vector(loss_fn(output))
+        losses = forward_pass(model_autogram, inputs, loss_fn, reduce_to_vector)
         _ = engine.compute_gramian(losses)
 
     # Verify that even with the hooked modules, autograd works normally when not using the engine.
     # Results should be the same as a normal call to autograd, and no time should be spent computing
     # the gramian at all.
-    torch.manual_seed(0)  # Fix randomness for random models
-    autograd_forward_backward(model_autogram, input, loss_fn)
+    autograd_forward_backward(model_autogram, inputs, loss_fn)
     grads = {name: p.grad for name, p in model_autogram.named_parameters() if p.grad is not None}
 
     assert_tensor_dicts_are_close(grads, autograd_grads)
@@ -398,7 +356,7 @@ def test_autograd_while_modules_are_hooked(
     ["factory", "batch_dim"],
     [
         (ModuleFactory(InstanceNorm2d, num_features=3, affine=True, track_running_stats=True), 0),
-        (ModuleFactory(RNN, input_size=8, hidden_size=5, batch_first=True), 0),
+        param(ModuleFactory(WithRNN), 0),
         (ModuleFactory(BatchNorm2d, num_features=3, affine=True, track_running_stats=False), 0),
     ],
 )
@@ -418,12 +376,11 @@ def test_compute_gramian_manual():
 
     in_dims = 18
     out_dims = 25
-
     factory = ModuleFactory(Linear, in_dims, out_dims)
     model = factory()
-    engine = Engine(model, batch_dim=None)
-
     input = randn_(in_dims)
+
+    engine = Engine(model, batch_dim=None)
     output = model(input)
     gramian = engine.compute_gramian(output)
 
@@ -464,21 +421,19 @@ def test_reshape_equivariance(shape: list[int]):
 
     input_size = shape[0]
     output_size = prod(shape[1:])
-
     factory = ModuleFactory(Linear, input_size, output_size)
     model1, model2 = factory(), factory()
+    input = randn_([input_size])
 
     engine1 = Engine(model1, batch_dim=None)
-    engine2 = Engine(model2, batch_dim=None)
-
-    input = randn_([input_size])
     output = model1(input)
-    reshaped_output = model2(input).reshape(shape[1:])
-
     gramian = engine1.compute_gramian(output)
-    reshaped_gramian = engine2.compute_gramian(reshaped_output)
     expected_reshaped_gramian = reshape_gramian(gramian, shape[1:])
 
+    engine2 = Engine(model2, batch_dim=None)
+    reshaped_output = model2(input).reshape(shape[1:])
+    reshaped_gramian = engine2.compute_gramian(reshaped_output)
+
     assert_close(reshaped_gramian, expected_reshaped_gramian)
 
 
@@ -504,21 +459,19 @@ def test_movedim_equivariance(shape: list[int], source: list[int], destination:
 
     input_size = shape[0]
     output_size = prod(shape[1:])
-
     factory = ModuleFactory(Linear, input_size, output_size)
     model1, model2 = factory(), factory()
+    input = randn_([input_size])
 
     engine1 = Engine(model1, batch_dim=None)
-    engine2 = Engine(model2, batch_dim=None)
-
-    input = randn_([input_size])
     output = model1(input).reshape(shape[1:])
-    moved_output = model2(input).reshape(shape[1:]).movedim(source, destination)
-
     gramian = engine1.compute_gramian(output)
-    moved_gramian = engine2.compute_gramian(moved_output)
     expected_moved_gramian = movedim_gramian(gramian, source, destination)
 
+    engine2 = Engine(model2, batch_dim=None)
+    moved_output = model2(input).reshape(shape[1:]).movedim(source, destination)
+    moved_gramian = engine2.compute_gramian(moved_output)
+
     assert_close(moved_gramian, expected_moved_gramian)
 
 
@@ -547,18 +500,16 @@ def test_batched_non_batched_equivalence(shape: list[int], batch_dim: int):
     input_size = prod(non_batched_shape)
     batch_size = shape[batch_dim]
     output_size = input_size
-
     factory = ModuleFactory(Linear, input_size, output_size)
     model1, model2 = factory(), factory()
+    input = randn_([batch_size, input_size])
 
     engine1 = Engine(model1, batch_dim=batch_dim)
-    engine2 = Engine(model2, batch_dim=None)
-
-    input = randn_([batch_size, input_size])
     output1 = model1(input).reshape([batch_size] + non_batched_shape).movedim(0, batch_dim)
-    output2 = model2(input).reshape([batch_size] + non_batched_shape).movedim(0, batch_dim)
-
     gramian1 = engine1.compute_gramian(output1)
+
+    engine2 = Engine(model2, batch_dim=None)
+    output2 = model2(input).reshape([batch_size] + non_batched_shape).movedim(0, batch_dim)
     gramian2 = engine2.compute_gramian(output2)
 
     assert_close(gramian1, gramian2)
@@ -575,24 +526,15 @@ def test_batched_non_batched_equivalence_2(factory: ModuleFactory, batch_size: i
     """
 
     model_0, model_none = factory(), factory()
-    input_shapes, output_shapes = get_in_out_shapes(model_0)
-
-    engine_0 = Engine(model_0, batch_dim=0)
-    engine_none = Engine(model_none, batch_dim=None)
-
-    inputs = make_tensors(batch_size, input_shapes)
-    targets = make_tensors(batch_size, output_shapes)
+    inputs, targets = make_inputs_and_targets(model_0, batch_size)
     loss_fn = make_mse_loss_fn(targets)
 
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_0(inputs)
-    losses_0 = reduce_to_vector(loss_fn(output))
-
-    torch.random.manual_seed(0)  # Fix randomness for random models
-    output = model_none(inputs)
-    losses_none = reduce_to_vector(loss_fn(output))
-
+    engine_0 = Engine(model_0, batch_dim=0)
+    losses_0 = forward_pass(model_0, inputs, loss_fn, reduce_to_vector)
     gramian_0 = engine_0.compute_gramian(losses_0)
+
+    engine_none = Engine(model_none, batch_dim=None)
+    losses_none = forward_pass(model_none, inputs, loss_fn, reduce_to_vector)
     gramian_none = engine_none.compute_gramian(losses_none)
 
     assert_close(gramian_0, gramian_none, rtol=1e-4, atol=1e-5)