diff --git a/pylops/torchoperator.py b/pylops/torchoperator.py index 7a6ad02e..58126bc9 100644 --- a/pylops/torchoperator.py +++ b/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) diff --git a/pytests/test_torchoperator.py b/pytests/test_torchoperator.py index a3690c33..81d19c85 100644 --- a/pytests/test_torchoperator.py +++ b/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) diff --git a/tutorials/torchop.py b/tutorials/torchop.py index 1d14de58..a5e1e42a 100755 --- a/tutorials/torchop.py +++ b/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)