#### 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).. # Distributed Training of Unfolded Networks In many large-scale imaging problems, the size of the image/volume to reconstruct is very large, making it impossible to train reconstruction networks (in this example, unfolded networks) with a single GPU. The `deepinv.distributed` framework enables training a model on multiple GPUs, by carefully parallelizing the data fidelity and denoising steps inside the network. This example shows how to combine: - the distributed framework (image/model parallelism over large images), - unfolded optimization with [`deepinv.optim.DRS`](https://deepinv.org/api/stubs/deepinv.optim.DRS.html.md#deepinv.optim.DRS), - standard training with [`deepinv.Trainer`](https://deepinv.org/api/stubs/deepinv.Trainer.html.md#deepinv.Trainer). Each GPU (rank) processes different parts/operators of the same image. This is not standard data-parallel training (e.g., via [`torch.nn.parallel.DistributedDataParallel`](https://docs.pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel)) over different images. **Usage:** ```bash # Single process python examples/distributed/demo_unrolled_distributed.py ``` ```bash # Multi-process (2 ranks) python -m torch.distributed.run --nproc_per_node=2 examples/distributed/demo_unrolled_distributed.py ``` ```Python # Import modules # ----------------------------------------------------------------------------- import os import torch from torch.utils.data import DataLoader, Subset from torchvision import transforms import deepinv as dinv from deepinv.datasets import HDF5Dataset, generate_dataset from deepinv.distributed import DistributedContext, distribute from deepinv.loss.metric import PSNR from deepinv.models import DRUNet from deepinv.optim import DRS from deepinv.optim.data_fidelity import L2 from deepinv.optim.prior import PnP from deepinv.physics import GaussianNoise, stack from deepinv.physics.blur import Blur from deepinv.physics.functional import gaussian_blur from deepinv.utils import get_data_home from deepinv.utils.plotting import plot, plot_curves from deepinv.utils.tensorlist import TensorList ``` # Dataset and Dataloader preparation helper functions In this example, we use the Urban100 dataset to generate measurements for training and validation. For every clean image, [`deepinv.datasets.generate_dataset()`](https://deepinv.org/api/stubs/deepinv.datasets.generate_dataset.html.md#deepinv.datasets.generate_dataset) creates two blurred and noisy measurements. Only rank 0 downloads and generates the dataset; the other ranks wait until both HDF5 files have been closed before opening them. ```Python def collate_batch(batch): """Collate clean/measured pairs while preserving TensorList measurements.""" if len(batch) == 1: x, y = batch[0] if x.ndim == 3: x = x.unsqueeze(0) y = TensorList([m.unsqueeze(0) if m.ndim == 3 else m for m in y]) return x, y xs = [x for x, _ in batch] ys = [y for _, y in batch] x_batch = torch.stack(xs, dim=0) n_ops = len(ys[0]) return x_batch, TensorList( [torch.stack([yy[i] for yy in ys], dim=0) for i in range(n_ops)] ) def prepare_dataset( ctx: DistributedContext, seed: int, crop_size: int, train_images: int, val_images: int, batch_size: int, num_workers: int, dataset_name: str, ): """Create/load Urban100 measurements for training and validation. Important: all ranks iterate over the same batches since distribution is over image content/operators, not over different images. """ blur_rngs = [ torch.Generator(device=ctx.device).manual_seed(seed + i) for i in range(2) ] physics_list = [ Blur( filter=gaussian_blur(sigma=(1.5, 1.5), device=str(ctx.device)), padding="circular", device=ctx.device, noise_model=GaussianNoise(sigma=0.03, rng=blur_rngs[0]), ), Blur( filter=gaussian_blur(sigma=(2.0, 2.0), device=str(ctx.device)), padding="circular", device=ctx.device, noise_model=GaussianNoise(sigma=0.05, rng=blur_rngs[1]), ), ] stacked_physics = stack(*physics_list) data_root = get_data_home() / "Urban100" os.makedirs(data_root, exist_ok=True) transform = transforms.Compose( [ transforms.Resize(crop_size), transforms.CenterCrop(crop_size), transforms.ToTensor(), ] ) if ctx.rank == 0: base_dataset = dinv.datasets.Urban100HR( root=str(data_root), download=True, transform=transform ) max_images = min(len(base_dataset), train_images + val_images) train_base = Subset(base_dataset, list(range(train_images))) val_base = Subset(base_dataset, list(range(train_images, max_images))) generate_dataset( train_dataset=train_base, physics=stacked_physics, save_dir=str(data_root), dataset_filename=f"{dataset_name}_train", device=ctx.device, train_datapoints=train_images, num_workers=num_workers, ) generate_dataset( train_dataset=val_base, physics=stacked_physics, save_dir=str(data_root), dataset_filename=f"{dataset_name}_val", device=ctx.device, train_datapoints=len(val_base), num_workers=num_workers, ) # generate_dataset closes each HDF5 file before returning. The data directory # must be on storage shared by all ranks. ctx.barrier() ensures that all ranks wait # until the HDF5 files are closed before opening them. ctx.barrier() train_ds = HDF5Dataset( path=str(data_root / f"{dataset_name}_train0.h5"), train=True ) val_ds = HDF5Dataset(path=str(data_root / f"{dataset_name}_val0.h5"), train=True) train_generator = torch.Generator().manual_seed(seed + 123) val_generator = torch.Generator().manual_seed(seed + 456) train_loader = DataLoader( train_ds, batch_size=batch_size, shuffle=True, generator=train_generator, num_workers=num_workers, collate_fn=collate_batch, ) val_loader = DataLoader( val_ds, batch_size=batch_size, shuffle=False, generator=val_generator, num_workers=num_workers, collate_fn=collate_batch, ) return stacked_physics, train_loader, val_loader ``` # Configuration Settings for training and distributed processing. patch_size and overlap control the size of the image patches that each rank processes, and how much they overlap with each other. #### NOTE The following settings are for demonstration purposes. We recommend training for more epochs to get better results. ```Python seed = 0 n_unroll = 3 # Number of unrolled iterations (DRS steps). crop_size = 128 if torch.cuda.is_available() else 64 # Training and dataloader settings epochs = 2 if torch.cuda.is_available() else 1 batch_size = 1 train_images = 16 if torch.cuda.is_available() else 6 val_images = 6 if torch.cuda.is_available() else 4 learning_rate = 2e-4 num_workers = 4 if torch.cuda.is_available() else 0 # Distributed processing settings patch_size = crop_size // 2 overlap = max(8, patch_size // 8) _ = torch.manual_seed(seed) ``` # Build distributed physics/model and train with deepinv.Trainer The distributed framework allows to distribute unfolded network with a few simple steps: - Initialize the distributed context - Prepare the physics, model, trainer and dataloaders - Call [`deepinv.distributed.distribute()`](https://deepinv.org/api/stubs/deepinv.distributed.distribute.html.md#deepinv.distributed.distribute) to distribute the physics and model across ranks - Train with [`deepinv.Trainer`](https://deepinv.org/api/stubs/deepinv.Trainer.html.md#deepinv.Trainer) as usual. The framework takes care of synchronizing the forward/backward passes across ranks, and communicating the necessary information between them. Reload checkpoints with the usual DeepInverse Trainer API. In this run, rank 0 can call `model = trainer.load_best_model()`. In a new script, rebuild the same model and Trainer, then call `trainer.load_model("ckpts/distributed_unfolded_drs//ckp_best.pth.tar")`. No rank-specific path or distributed checkpoint API is needed. ```Python # Keep identical random streams across ranks: this framework splits each image # across devices, so all ranks should consume the same minibatches. with DistributedContext(seed=seed, seed_offset=False) as ctx: if ctx.rank == 0: print(f"Processes: {ctx.world_size}") print(f"Device: {ctx.device}") stacked_physics, train_loader, val_loader = prepare_dataset( ctx, seed=seed, crop_size=crop_size, train_images=train_images, val_images=val_images, batch_size=batch_size, num_workers=num_workers, dataset_name="urban100_drs_blur_noise", ) # Distribute the stacked physics across ranks. distributed_physics = distribute( stacked_physics, ctx, ) # Build an unfolded DRS model and distribute trainable components. denoiser = DRUNet(pretrained="download").to(ctx.device) prior = PnP(denoiser=denoiser) model = DRS( stepsize=[0.9] * n_unroll, sigma_denoiser=[0.04] * n_unroll, beta=[1.0] * n_unroll, trainable_params=["stepsize", "sigma_denoiser", "beta"], data_fidelity=L2(), prior=prior, max_iter=n_unroll, unfold=True, ) model = distribute( model, ctx, patch_size=patch_size, overlap=overlap, max_batch_size=1, ) optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) psnr_metric = PSNR(reduction="mean") # The trainable model parameters are synchronized across ranks, so every # rank holds equivalent weights. Save only one representative checkpoint. checkpoint_root = "ckpts/distributed_unfolded_drs" if ctx.rank == 0 else None # Reconstruction before training. demo_x, demo_y = next(iter(val_loader)) demo_x = demo_x.to(ctx.device) demo_y = demo_y.to(ctx.device) with torch.no_grad(): demo_rec_before = model(demo_y, distributed_physics) trainer = dinv.Trainer( model=model, physics=distributed_physics, epochs=epochs, device=ctx.device, losses=[dinv.loss.SupLoss(metric=dinv.metric.MSE())], metrics=psnr_metric, optimizer=optimizer, train_dataloader=train_loader, eval_dataloader=val_loader, grad_clip=1.0, compare_no_learning=False, save_path=checkpoint_root, verbose=(ctx.rank == 0), show_progress_bar=(ctx.rank == 0), freq_update_progress_bar=5, check_grad=True, non_blocking_transfers=False, ) trainer.train() with torch.no_grad(): demo_rec_after = model(demo_y, distributed_physics) # Display training summary and qualitative result (rank 0 only) if ctx.rank == 0: train_history = trainer.train_metrics_history.get("PSNR", []) val_history = trainer.eval_metrics_history.get("PSNR", []) final_steps = [f"{p.item():.4f}" for p in model.params_algo["stepsize"]] print(f"Final trainable stepsizes: {final_steps}") if val_history: print(f"Final val PSNR: {val_history[-1]:.2f} dB") plot( [demo_x, demo_y[0], demo_rec_before, demo_rec_after], titles=[ "Ground truth", "Blurred noisy measurement", "Before training", "After training", ], save_fn="distributed_unrolled_result.png", ) if train_history and val_history and len(train_history) > 1: plot_curves({"train_psnr": [train_history], "val_psnr": [val_history]}) print("Saved: distributed_unrolled_result.png") ``` ```none Processes: 1 Device: cuda:0 /local/jtachell/deepinv/deepinv/examples/distributed/demo_unrolled_distributed.py:120: DeprecationWarning: Function 'get_data_home' is deprecated and will be removed in a future version. data_root = get_data_home() / "Urban100" 0%| | 0/135388067 [00:00 [`Download Jupyter notebook: demo_unrolled_distributed.ipynb`](https://deepinv.org/auto_examples/distributed/_downloads/fcbf8f36ff078d902cb51b06a3c17ceb/demo_unrolled_distributed.ipynb) [`Download Python source code: demo_unrolled_distributed.py`](https://deepinv.org/auto_examples/distributed/_downloads/5d7399625640718d6cf2eb286e0f4c07/demo_unrolled_distributed.py) [`Download zipped: demo_unrolled_distributed.zip`](https://deepinv.org/auto_examples/distributed/_downloads/03595bf9c6d20bba31513d45edbd4440/demo_unrolled_distributed.zip) [Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)