Merge pull request #750 from mrava87/feat-torchopadjoint

Feat: added forward/adjoint to TorchOperator

Author: mrava87

pylops/torchoperator.py

@@ -3,6 +3,8 @@
 ]
 
 
+from math import prod
+
 import numpy as np
 
 from pylops import LinearOperator
@@ -91,6 +93,15 @@ def __init__(
     def __call__(self, x):
         return self.apply(x)
 
+    def __repr__(self):
+        M, N = prod(self.dimsd), prod(self.dims)
+        if self.dtype is None:
+            dt = "unspecified dtype"
+        else:
+            dt = "dtype=" + str(self.dtype)
+
+        return "<%dx%d %s with %s>" % (M, N, self.__class__.__name__, dt)
+
     def apply(self, x: TensorTypeLike) -> TensorTypeLike:
         """Apply forward pass to input vector
 
@@ -106,3 +117,22 @@ def apply(self, x: TensorTypeLike) -> TensorTypeLike:
 
         """
         return self.Top(x, self.matvec, self.rmatvec, self.device, self.devicetorch)
+
+    # alias for forward pass
+    forward = apply
+
+    def adjoint(self, x: TensorTypeLike) -> TensorTypeLike:
+        """Apply adjoint pass to input vector
+
+        Parameters
+        ----------
+        x : :obj:`torch.Tensor`
+            Input array
+
+        Returns
+        -------
+        y : :obj:`torch.Tensor`
+            Output array resulting from the application of the adjoint operator to ``x``.
+
+        """
+        return self.Top(x, self.rmatvec, self.matvec, self.device, self.devicetorch)

pytests/test_torchoperator.py

@@ -109,3 +109,40 @@ def test_TorchOperator_batch_nd(par, dtype):
 
     assert yt.dtype == dtype
     assert_array_equal(y, yt)
+
+
+@pytest.mark.skipif(platform.system() == "Darwin", reason="Not OSX enabled")
+@pytest.mark.parametrize("par", [(par1)])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_TorchOperator_forward_adjoint(par, dtype):
+    """Compute gradient of L2 norm (chain of forward and adjoint) and
+    compare Jacobian vector product with analytical solution"""
+    device = "cpu" if backend == "numpy" else "cuda"
+
+    Dop = MatrixMult(
+        np.random.normal(0.0, 1.0, (par["ny"], par["nx"])).astype(dtype), dtype=dtype
+    )
+    Top = TorchOperator(Dop, batch=False, device="cpu" if backend == "numpy" else "gpu")
+
+    x = np.random.normal(0.0, 1.0, par["nx"]).astype(dtype)
+    xt = torch.from_numpy(to_numpy(x)).to(device).view(-1)
+    xt.requires_grad = True
+    y = -2 * np.arange(par["ny"], dtype=dtype)
+    yt = torch.from_numpy(to_numpy(y)).to(device).view(-1)
+    v = np.random.normal(0.0, 1.0, par["ny"]).astype(dtype)
+    vt = torch.from_numpy(to_numpy(v)).to(device).view(-1)
+
+    # pylops operator
+    f = Dop.H * (y - Dop * x)
+    jvt = -Dop.H * Dop * v
+
+    # torch operator
+    ft = Top.adjoint(yt - Top.forward(xt))
+    ft.backward(vt, retain_graph=True)
+    jvtt = xt.grad.cpu().numpy()
+    ft = ft.detach().cpu().numpy()
+
+    assert ft.dtype == x.dtype
+    assert jvtt.dtype == x.dtype
+    assert_array_equal(f, ft)
+    assert_array_equal(jvt, jvtt)

tutorials/torchop.py

@@ -149,7 +149,13 @@ def forward(self, x):
 plt.tight_layout()
 
 ###############################################################################
-# And finally we do the same with a batch of 3 training samples.
+# And we can do the same with a batch of 3 training samples. Note that under
+# the hood, this effectively calls the matrix-matrix version of the forward
+# and adjoint operator (i.e., `matmat` and `rmatmat`); for operators that do
+# not implement these methods directly, this is simply implemented by calling
+# the matrix-vector of the forward and adjoint operator (i.e., `matvec` and
+# `rmatvec`)multiple times, which is less efficient.
+
 net = Network(4)
 Cop = pylops.TorchOperator(pylops.Smoothing2D((5, 5), dims=(32, 32)), batch=True)
 
@@ -169,3 +175,53 @@ def forward(self, x):
 axs[1].set_title("Gradient")
 axs[1].axis("tight")
 plt.tight_layout()
+
+###############################################################################
+# Finally, whilst :class:`pylops.TorchOperator` is designed such that
+# when a PyLops linear operator is inserted into a Torch graph, the backward
+# pass will automatically call the adjoint of the operator, it is also possible to
+# explicitly call the forward and adjoint of the operator in the forward pass of
+# an AD chain. This can be useful in some scenarios, for example in the
+# implementation of so-called unrolled networks. In this case, we can simply
+# use the ``forward`` and ``adjoint`` methods of the :class:`pylops.TorchOperator`
+# class; Torch's AD will instead call the two methods swapped, namely ``adjoint``
+# and ``forward``.
+#
+# Let's consider the following example:
+#
+#   .. math::
+#        \mathbf{y}=\textbf{A}^H (\textbf{A} \mathbf{x} - \mathbf{d})
+#
+# whose Jacobian is given by:
+#
+#   .. math::
+#        \mathbf{J}=-\textbf{A}^H \textbf{A}
+#
+# Let's once again verify that the result of the product between
+# the transposed Jacobian and a vector :math:`\mathbf{v}` matches
+# with the analytical one.
+
+nx, ny = 10, 6
+xt0 = torch.arange(nx, dtype=torch.double, requires_grad=True)
+x0 = xt0.detach().numpy()
+yt0 = -2 * torch.arange(ny, dtype=torch.double)
+y0 = xt0.detach().numpy()
+
+# Forward
+A = np.random.normal(0.0, 1.0, (ny, nx))
+At = torch.from_numpy(A)
+Atop = pylops.TorchOperator(pylops.MatrixMult(A))
+yt = Atop.adjoint(yt0 - Atop.forward(xt0))
+
+# AD
+v = torch.ones(nx, dtype=torch.double)
+yt.backward(v, retain_graph=True)
+adgrad = xt0.grad
+
+# Analytical
+JT = -At.T @ At
+anagrad = torch.matmul(JT, v)
+
+print("Input: ", x0)
+print("AD gradient: ", adgrad)
+print("Analytical gradient: ", anagrad)