Merge branch 'main' into revamp-interface

Author: ValerianRey

README.md

@@ -111,11 +111,11 @@ using [UPGrad](https://torchjd.org/stable/docs/aggregation/upgrad/).
       loss1 = loss_fn(output1, target1)
       loss2 = loss_fn(output2, target2)
 
-      optimizer.zero_grad()
 -     loss = loss1 + loss2
 -     loss.backward()
 +     mtl_backward(losses=[loss1, loss2], features=features, aggregator=aggregator)
       optimizer.step()
+      optimizer.zero_grad()
 ```
 
 > [!NOTE]
@@ -150,12 +150,12 @@ Jacobian descent using [UPGrad](https://torchjd.org/stable/docs/aggregation/upgr
 -     loss = loss_fn(output, target)  # shape [1]
 +     losses = loss_fn(output, target)  # shape [16]
 
-      optimizer.zero_grad()
 -     loss.backward()
 +     gramian = engine.compute_gramian(losses)  # shape: [16, 16]
 +     weights = weighting(gramian)  # shape: [16]
 +     losses.backward(weights)
       optimizer.step()
+      optimizer.zero_grad()
 ```
 
 Lastly, you can even combine the two approaches by considering multiple tasks and each element of
@@ -201,10 +201,10 @@ for input, target1, target2 in zip(inputs, task1_targets, task2_targets):
     # Obtain the weights that lead to no conflict between reweighted gradients
     weights = weighting(gramian)  # shape: [16, 2]
 
-    optimizer.zero_grad()
     # Do the standard backward pass, but weighted using the obtained weights
     losses.backward(weights)
     optimizer.step()
+    optimizer.zero_grad()
 ```
 
 > [!NOTE]

docs/source/examples/amp.rst

@@ -54,6 +54,7 @@ following example shows the resulting code for a multi-task learning use-case.
         scaler.step(optimizer)
         scaler.update()
         optimizer.zero_grad()
+        optimizer.zero_grad()
 
 .. hint::
     Within the ``torch.autocast`` context, some operations may be done in ``float16`` type. For

docs/source/examples/iwrm.rst

@@ -70,6 +70,7 @@ batch of data. When minimizing per-instance losses (IWRM), we use either autojac
 
                 optimizer.step()
                 optimizer.zero_grad()
+                optimizer.zero_grad()
 
         In this baseline example, the update may negatively affect the loss of some elements of the
         batch.
@@ -106,6 +107,7 @@ batch of data. When minimizing per-instance losses (IWRM), we use either autojac
 
                 optimizer.step()
                 optimizer.zero_grad()
+                optimizer.zero_grad()
 
         Here, we compute the Jacobian of the per-sample losses with respect to the model parameters
         and use it to update the model such that no loss from the batch is (locally) increased.
@@ -142,6 +144,7 @@ batch of data. When minimizing per-instance losses (IWRM), we use either autojac
                 losses.backward(weights)
                 optimizer.step()
                 optimizer.zero_grad()
+                optimizer.zero_grad()
 
         Here, the per-sample gradients are never fully stored in memory, leading to large
         improvements in memory usage and speed compared to autojac, in most practical cases. The

src/torchjd/autogram/_engine.py

@@ -77,7 +77,7 @@ class Engine:
         Train a model using Gramian-based Jacobian descent.
 
         .. code-block:: python
-            :emphasize-lines: 5-6, 15-16, 18-19, 26-28
+            :emphasize-lines: 5-6, 15-16, 18-19, 26-29
 
             import torch
             from torch.nn import Linear, MSELoss, ReLU, Sequential
@@ -103,11 +103,11 @@ class Engine:
                 output = model(input).squeeze(dim=1)  # shape: [16]
                 losses = criterion(output, target)  # shape: [16]
 
-                optimizer.zero_grad()
                 gramian = engine.compute_gramian(losses)  # shape: [16, 16]
                 weights = weighting(gramian)  # shape: [16]
                 losses.backward(weights)
                 optimizer.step()
+                optimizer.zero_grad()
 
         This is equivalent to just calling ``torchjd.autojac.backward(losses, UPGrad())``. However,
         since the Jacobian never has to be entirely in memory, it is often much more

tests/doc/test_autogram.py

@@ -26,8 +26,8 @@ def test_engine():
         output = model(input).squeeze(dim=1)  # shape: [16]
         losses = criterion(output, target)  # shape: [16]
 
-        optimizer.zero_grad()
         gramian = engine.compute_gramian(losses)  # shape: [16, 16]
         weights = weighting(gramian)  # shape: [16]
         losses.backward(weights)
         optimizer.step()
+        optimizer.zero_grad()

tests/doc/test_rst.py

@@ -48,6 +48,7 @@ def test_amp():
         scaler.step(optimizer)
         scaler.update()
         optimizer.zero_grad()
+        optimizer.zero_grad()
 
 
 def test_basic_usage():
@@ -122,6 +123,7 @@ def test_iwmtl():
         losses.backward(weights)
         optimizer.step()
         optimizer.zero_grad()
+        optimizer.zero_grad()
 
 
 def test_iwrm():
@@ -145,6 +147,7 @@ def test_autograd():
             loss.backward()
             optimizer.step()
             optimizer.zero_grad()
+            optimizer.zero_grad()
 
     def test_autojac():
         import torch
@@ -200,6 +203,7 @@ def test_autogram():
             losses.backward(weights)
             optimizer.step()
             optimizer.zero_grad()
+            optimizer.zero_grad()
 
     test_autograd()
     test_autojac()
@@ -400,6 +404,7 @@ def test_partial_jd():
         losses.backward(weights)
         optimizer.step()
         optimizer.zero_grad()
+        optimizer.zero_grad()
 
 
 def test_rnn():

tests/unit/autogram/test_engine.py

@@ -345,7 +345,7 @@ def test_iwrm_steps_with_autogram(factory: ModuleFactory, batch_size: int, batch
         loss_fn = make_mse_loss_fn(targets)
         autogram_forward_backward(model, inputs, loss_fn, engine, weighting)
         optimizer.step()
-        model.zero_grad()
+        optimizer.zero_grad()
 
 
 @mark.parametrize(["factory", "batch_size"], PARAMETRIZATIONS)