Merge branch 'main' into linear-gramian-computer

Author: ValerianRey

CONTRIBUTING.md

@@ -130,18 +130,19 @@ from utils.tensors import ones_
 a = ones_(3, 4)
 ```
 
-This will automatically call `torch.ones` with `device=unit.conftest.DEVICE`.
+This will automatically call `torch.ones` with `device=DEVICE`.
 If the function you need does not exist yet as a partial function in `tensors.py`, add it.
 Lastly, when you create a model or a random generator, you have to move them manually to the right
-device (the `DEVICE` defined in `unit.conftest`):
+device (the `DEVICE` defined in `device.py`).
 ```python
 import torch
 from torch.nn import Linear
-from unit.conftest import DEVICE
+from device import DEVICE
 
 model = Linear(3, 4).to(device=DEVICE)
 rng = torch.Generator(device=DEVICE)
 ```
+You may also use a `ModuleFactory` to make the modules on `DEVICE` automatically.
 
 ### Coding
 

tests/unit/conftest.py renamed to tests/conftest.py

@@ -1,21 +1,13 @@
-import os
 import random as rand
+from contextlib import nullcontext
 
 import torch
-from pytest import fixture, mark
+from device import DEVICE
+from pytest import RaisesExc, fixture, mark
+from torch import Tensor
+from utils.architectures import ModuleFactory
 
-try:
-    _device_str = os.environ["PYTEST_TORCH_DEVICE"]
-except KeyError:
-    _device_str = "cpu"  # Default to cpu if environment variable not set
-
-if _device_str != "cuda:0" and _device_str != "cpu":
-    raise ValueError(f"Invalid value of environment variable PYTEST_TORCH_DEVICE: {_device_str}")
-
-if _device_str == "cuda:0" and not torch.cuda.is_available():
-    raise ValueError('Requested device "cuda:0" but cuda is not available.')
-
-DEVICE = torch.device(_device_str)
+from torchjd.aggregation import Aggregator, Weighting
 
 
 @fixture(autouse=True)
@@ -48,3 +40,24 @@ def pytest_collection_modifyitems(config, items):
             item.add_marker(skip_slow)
         if "xfail_if_cuda" in item.keywords and str(DEVICE).startswith("cuda"):
             item.add_marker(xfail_cuda)
+
+
+def pytest_make_parametrize_id(config, val, argname):
+    MAX_SIZE = 40
+    optional_string = None  # Returning None means using pytest's way of making the string
+
+    if isinstance(val, (Aggregator, ModuleFactory, Weighting)):
+        optional_string = str(val)
+    elif isinstance(val, Tensor):
+        optional_string = "T" + str(list(val.shape))  # T to indicate that it's a tensor
+    elif isinstance(val, (tuple, list, set)) and len(val) < 20:
+        optional_string = str(val)
+    elif isinstance(val, RaisesExc):
+        optional_string = " or ".join([f"{exc.__name__}" for exc in val.expected_exceptions])
+    elif isinstance(val, nullcontext):
+        optional_string = "does_not_raise()"
+
+    if isinstance(optional_string, str) and len(optional_string) > MAX_SIZE:
+        optional_string = optional_string[: MAX_SIZE - 3] + "+++"  # Can't use dots with pytest
+
+    return optional_string

tests/device.py

@@ -0,0 +1,16 @@
+import os
+
+import torch
+
+try:
+    _device_str = os.environ["PYTEST_TORCH_DEVICE"]
+except KeyError:
+    _device_str = "cpu"  # Default to cpu if environment variable not set
+
+if _device_str != "cuda:0" and _device_str != "cpu":
+    raise ValueError(f"Invalid value of environment variable PYTEST_TORCH_DEVICE: {_device_str}")
+
+if _device_str == "cuda:0" and not torch.cuda.is_available():
+    raise ValueError('Requested device "cuda:0" but cuda is not available.')
+
+DEVICE = torch.device(_device_str)

tests/speed/autogram/grad_vs_jac_vs_gram.py

@@ -2,18 +2,19 @@
 import time
 
 import torch
-from unit.conftest import DEVICE
+from device import DEVICE
 from utils.architectures import (
     AlexNet,
     Cifar10Model,
     FreeParam,
     GroupNormMobileNetV3Small,
     InstanceNormMobileNetV2,
     InstanceNormResNet18,
+    ModuleFactory,
     NoFreeParam,
-    ShapedModule,
     SqueezeNet,
     WithTransformerLarge,
+    get_in_out_shapes,
 )
 from utils.forward_backwards import (
     autograd_forward_backward,
@@ -28,33 +29,30 @@
 from torchjd.autogram import Engine
 
 PARAMETRIZATIONS = [
-    (WithTransformerLarge, 8),
-    (FreeParam, 64),
-    (NoFreeParam, 64),
-    (Cifar10Model, 64),
-    (AlexNet, 8),
-    (InstanceNormResNet18, 16),
-    (GroupNormMobileNetV3Small, 16),
-    (SqueezeNet, 4),
-    (InstanceNormMobileNetV2, 2),
+    (ModuleFactory(WithTransformerLarge), 8),
+    (ModuleFactory(FreeParam), 64),
+    (ModuleFactory(NoFreeParam), 64),
+    (ModuleFactory(Cifar10Model), 64),
+    (ModuleFactory(AlexNet), 8),
+    (ModuleFactory(InstanceNormResNet18), 16),
+    (ModuleFactory(GroupNormMobileNetV3Small), 16),
+    (ModuleFactory(SqueezeNet), 4),
+    (ModuleFactory(InstanceNormMobileNetV2), 2),
 ]
 
 
-def compare_autograd_autojac_and_autogram_speed(architecture: type[ShapedModule], batch_size: int):
-    input_shapes = architecture.INPUT_SHAPES
-    output_shapes = architecture.OUTPUT_SHAPES
+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)
     loss_fn = make_mse_loss_fn(targets)
 
-    model = architecture().to(device=DEVICE)
-
     A = Mean()
     W = A.weighting
 
     print(
-        f"\nTimes for forward + backward on {architecture.__name__} with BS={batch_size}, A={A}"
-        f" on {DEVICE}."
+        f"\nTimes for forward + backward on {factory} with BS={batch_size}, A={A}" f" on {DEVICE}."
     )
 
     def fn_autograd():
@@ -148,8 +146,8 @@ def time_call(fn, init_fn=noop, pre_fn=noop, post_fn=noop, n_runs: int = 10) ->
 
 
 def main():
-    for architecture, batch_size in PARAMETRIZATIONS:
-        compare_autograd_autojac_and_autogram_speed(architecture, batch_size)
+    for factory, batch_size in PARAMETRIZATIONS:
+        compare_autograd_autojac_and_autogram_speed(factory, batch_size)
         print("\n")
 
 

tests/unit/aggregation/_inputs.py

@@ -1,5 +1,5 @@
 import torch
-from unit.conftest import DEVICE
+from device import DEVICE
 from utils.tensors import zeros_
 
 from ._matrix_samplers import NonWeakSampler, NormalSampler, StrictlyWeakSampler, StrongSampler