#### NOTE New to DeepInverse? Get started with the basics with the [5 minute quickstart tutorial](https://deepinv.org/auto_examples/basics/demo_quickstart.html.md#sphx-glr-auto-examples-basics-demo-quickstart-py).. # Reducing the memory and computational complexity of unfolded network training Some unfolded architectures rely on a [`least-squares solver`](https://deepinv.org/api/stubs/deepinv.optim.linear.least_squares.html.md#deepinv.optim.linear.least_squares) to compute the proximal step w.r.t. the data-fidelity term (e.g., [`ADMM`](https://deepinv.org/api/stubs/deepinv.optim.optim_iterators.ADMMIteration.html.md#deepinv.optim.optim_iterators.ADMMIteration) or [`HQS`](https://deepinv.org/api/stubs/deepinv.optim.optim_iterators.HQSIteration.html.md#deepinv.optim.optim_iterators.HQSIteration)): $$ \operatorname{prox}_{\gamma f}(z) = \underset{x}{\arg\min} \; \frac{\gamma}{2}\|A_\theta x-y\|^2 + \frac{1}{2}\|x-z\|^2 $$ During backpropagation, a naive implementation requires storing the gradients of every intermediate step of the least squares solver (which is an iterative method), resulting in significant memory and computational costs which are proportional to number of iterations done by the solver. The library provides a memory-efficient back-propagation strategy that reduces the memory footprint during training, by computing the gradients of the proximal step in closed-form, without storing any intermediate steps. This closed-form computation requires evaluating the least-squares solver one additional time during the gradient computation. Let $h(z, y, \theta, \gamma) = \operatorname{prox}_{\gamma f}(z)$ be the proximal operator. During the backward pass, we need to compute the vector-Jacobian products (VJPs), w.r.t the input variables $(z, y, \theta, \gamma)$ required for backpropagation: $$ \left( \frac{\partial h}{\partial z} \right)^{\top} v, \quad \left( \frac{\partial h}{\partial y} \right)^{\top} v, \quad \left( \frac{\partial h}{\partial \theta} \right)^{\top} v, \quad \left( \frac{\partial h}{\partial \gamma} \right)^{\top} v $$ and $v$ is the upstream gradient. The VJPs can be computed in closed-form by solving a second least-squares problem, as shown in the following. When the forward least-squares solver converges to the exact minimizer, we have the following closed-form expressions for $h(z, y, \theta, \gamma)$: $$ h(z, y, \theta, \gamma) = \left( A_\theta^{\top} A_\theta + \frac{1}{\gamma} I \right)^{-1} \left( A_\theta^{\top} y + \frac{1}{\gamma} z \right) $$ Let $M$ denote the inverse $\left( A_\theta^T A_\theta + \frac{1}{\gamma} I \right)^{-1}$. The VJPs can be computed as follows: $$ \left( \frac{\partial h}{\partial z} \right)^{\top} v &= \frac{1}{\gamma} M v \\ \left( \frac{\partial h}{\partial y} \right)^{\top} v &= A_\theta M v \\ \left( \frac{\partial h}{\partial \gamma} \right)^{\top} v &= (h - z)^\top M v / \gamma^2 \\ \left( \frac{\partial h}{\partial \theta} \right)^{\top} v &= \frac{\partial p}{\partial \theta} $$ where $p = (y - A_\theta h)^{\top} A_\theta M v$ and $\frac{\partial p}{\partial \theta}$ can be computed using the standard backpropagation mechanism (autograd). #### NOTE Linear forward operators that have a [`closed-form singular value decomposition`](https://deepinv.org/api/stubs/deepinv.physics.DecomposablePhysics.html.md#deepinv.physics.DecomposablePhysics) benefit from a [`closed-form formula`](https://deepinv.org/api/stubs/deepinv.physics.DecomposablePhysics.html.md#deepinv.physics.DecomposablePhysics.prox_l2) for computing the proximal step, and thus we shouldn’t expect speed-ups in these specific cases. This example shows how to train an unfolded neural network with a memory complexity that is independent of the number of iterations in least squares solver (O(1) memory complexity) used for computing the data-fidelity proximal step. #### NOTE By default, this example is run on CPU so we should not expect significant speed-ups and we can not trace the memory usage precisely. For a better experience, we recommend running the example on a machine with a GPU. In such a case, we can expect significant speed-ups (20-50%) and a significant reduction in memory usage (2x-3x reduction factor). ```Python import deepinv as dinv import torch from torch.utils.data import DataLoader from deepinv.optim.data_fidelity import L2 from deepinv.optim.prior import PnP from torchvision import transforms from deepinv.utils.demo import load_dataset import time import numpy as np import matplotlib.pyplot as plt from deepinv.optim import HQS device = dinv.utils.get_device() dtype = torch.float32 img_size = 64 if torch.cuda.is_available() else 32 num_images = 64 ``` ```none Selected GPU 0 with 25271.25 MiB free memory ``` ## Define the degradation operator and the dataset. We use the Blur class with `valid` padding from the physics module. ITs proximal operator does not have a closed-form solution, and thus requires using a least-squares solver. ```Python # In this example, we use the CBSD500 dataset train_dataset_name = "CBSD500" # Specify the transforms to be applied to the input images. train_transform = transforms.Compose( [transforms.RandomCrop(img_size), transforms.ToTensor()] ) # Define the base train dataset of clean images. train_dataset = load_dataset(train_dataset_name, transform=train_transform) train_dataset = torch.utils.data.Subset(train_dataset, range(num_images)) train_loader = DataLoader( train_dataset, batch_size=32 if torch.cuda.is_available() else 8, num_workers=4 if torch.cuda.is_available() else 0, shuffle=True, ) physics = dinv.physics.Blur( filter=dinv.physics.functional.gaussian_blur(sigma=(2.5, 2.5)), padding="valid", device=device, noise_model=dinv.physics.GaussianNoise(sigma=0.1), max_iter=100, tol=1e-8, implicit_backward_solver=False, ) ``` ```none /local/jtachell/deepinv/deepinv/deepinv/datasets/base.py:495: DeprecationWarning: The tuple format for dataset outputs is deprecated and will be removed in a future version.It is recommended to set `use_dict_output=True` for better readability and flexibility in returned outputs.The default is currently `False` for backward compatibility, but will be switched to `True` in a future version. super().__init__(use_dict_output=use_dict_output) ``` ## Define the unfolded parameters. ```Python # Unrolled optimization algorithm parameters max_iter = 5 # number of unfolded layers # Select the data fidelity term data_fidelity = L2() stepsize = [1] * max_iter # stepsize of the algorithm sigma_denoiser = [0.01] * max_iter # noise level parameter of the denoiser trainable_params = [ "sigma_denoiser", "stepsize", ] # define which parameters are trainable ``` ## Train the network Here we write explicitly the training loop to show how implicit differentiation can be used to avoid out-of-memory issues and sometimes accelerate training. But you can also use the [`deepinv.Trainer`](https://deepinv.org/api/stubs/deepinv.Trainer.html.md#deepinv.Trainer) class as shown in other examples. ```Python # Some helper functions for measuring memory usage use_cuda = device.type == "cuda" and torch.cuda.is_available() def sync(): if use_cuda: torch.cuda.synchronize() def reset_memory(): if use_cuda: torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() def peak_memory(): if use_cuda: peak_bytes = int(torch.cuda.max_memory_allocated(device=device)) else: # Stats on CPU is not very accurate peak_bytes = 0 return peak_bytes def train_and_measure(): torch.manual_seed(42) prior = PnP(denoiser=dinv.models.DnCNN(depth=3, pretrained=None).to(device)) model = HQS( unfold=True, stepsize=stepsize, sigma_denoiser=sigma_denoiser, trainable_params=trainable_params, data_fidelity=data_fidelity, max_iter=max_iter, prior=prior, ).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=5e-4, weight_decay=1e-8) model.train() reset_memory() sync() start = time.perf_counter() losses = [] for x in train_loader: x = x.to(device) y = physics(x) optimizer.zero_grad() x_hat = model(physics=physics, y=y) loss = torch.nn.functional.mse_loss(x_hat, x) losses.append(loss.item()) loss.backward() optimizer.step() sync() end = time.perf_counter() peak_memory_mb = peak_memory() / (10**6) time_per_iter = (end - start) / len(train_loader) avg_loss = np.array(losses) avg_loss = np.cumsum(avg_loss) / (np.arange(len(avg_loss)) + 1) return dict( avg_loss=avg_loss, time_per_iter=time_per_iter, peak_memory_mb=peak_memory_mb ) ``` We first train the model will full backpropagation to compare the memory usage. Define the unfolded trainable model. ```Python # Setting this parameter to False to use full backpropagation physics.implicit_backward_solver = False auto = train_and_measure() ``` ```none /local/jtachell/deepinv/deepinv/.pixi/envs/docs/lib/python3.12/multiprocessing/popen_fork.py:66: DeprecationWarning: This process (pid=4088353) is multi-threaded, use of fork() may lead to deadlocks in the child. self.pid = os.fork() ``` We now train the model using the closed-form gradients of the proximal step. We can do this by setting `implicit_backward_solver` to `True`. ```Python physics.implicit_backward_solver = True implicit = train_and_measure() ``` ## Compare the memory usage ```Python print(f"Full backpropagation: time per iteration: {auto['time_per_iter']:.2f} s. ") print( f"Implicit differentiation: time per iteration: {implicit['time_per_iter']:.2f} s." ) # Compare the memory usage if use_cuda: print(f"Full backpropagation: peak memory usage: {auto['peak_memory_mb']:.1f} MB") print( f"Implicit differentiation: peak memory usage: {implicit['peak_memory_mb']:.1f} MB" ) print( f"Memory reduction factor: {auto['peak_memory_mb']/implicit['peak_memory_mb']:.1f}x" ) # Compare the training loss plt.figure(figsize=(7, 4)) plt.plot(auto["avg_loss"], label="Full backpropagation", linestyle="--", linewidth=2) plt.plot( implicit["avg_loss"], label="Implicit differentiation", linestyle="-.", linewidth=1.5, ) plt.yscale("log") plt.xlabel("Iteration", fontsize=12) plt.ylabel("Training loss (MSE)", fontsize=12) plt.legend() plt.title( f"Training loss. Avg loss difference: {np.mean(np.abs(auto['avg_loss'] - implicit['avg_loss'])):.2e}", fontsize=14, ) plt.grid() plt.show() ``` ```none Full backpropagation: time per iteration: 22.09 s. Implicit differentiation: time per iteration: 42.25 s. Full backpropagation: peak memory usage: 4199.0 MB Implicit differentiation: peak memory usage: 1239.2 MB Memory reduction factor: 3.4x ``` **Total running time of the script:** (2 minutes 9.468 seconds) [`Download Jupyter notebook: demo_unfolded_constant_memory.ipynb`](https://deepinv.org/auto_examples/unfolded/_downloads/c7e36bcaa21741ac5318f37ab8658137/demo_unfolded_constant_memory.ipynb) [`Download Python source code: demo_unfolded_constant_memory.py`](https://deepinv.org/auto_examples/unfolded/_downloads/0de4d34cd0bb27e15e0b678736214c70/demo_unfolded_constant_memory.py) [`Download zipped: demo_unfolded_constant_memory.zip`](https://deepinv.org/auto_examples/unfolded/_downloads/678fb0c5d7ffaa7d6b8f3da670b1f0b9/demo_unfolded_constant_memory.zip) [Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)