Plug-and-Play Reconstruction#

Usage and Principle#

The inference benchmark exposes this solver as PnP. One BenchOpt callback step performs one plug-and-play proximal-gradient iteration. Starting from \(x_k\), the solver first takes an L2 data-fidelity gradient step and then applies DRUNet as an implicit image prior:

\[\begin{split}z_k &= x_k - \gamma \nabla f(x_k), \\ x_{k+1} &= D_{\sigma}(z_k).\end{split}\]

The result is clipped to the dataset signal range by default. An optional relaxation blends the gradient-step result and denoised result instead of replacing it directly. The objective evaluates the current reconstruction after each callback, so max-runs in the experiment YAML controls the number of PnP iterations.

A minimal configuration is:

solver:
  - PnP:
      denoiser: drunet
      denoiser_sigma: 0.01
      step_size: 0.1
      init_method: zeros
      profiler_mode: custom

Multi-GPU Distribution#

Distributed execution uses one process per rank and ordinarily one GPU per process. The process topology comes from SLURM or torchrun; a DeepInv distributed context assigns the local device and coordinates collectives. Every rank follows the same callback decision and receives the assembled iterate.

The two expensive stages can be distributed independently:

distribute_physics: Partitions the stacked acquisition operators across ranks. Each rank computes its local forward/adjoint contribution and the data-fidelity result is reduced across ranks. This is most useful when there are several acquisition operators, views, or frames.
distribute_denoiser: Tiles the full image or volume into overlapping patches, assigns patches to ranks, evaluates DRUNet locally, and combines them with overlap-aware blending and an all-reduce. For useful scaling, create at least as many patches as ranks. patch_size, overlap, and max_batch_size trade communication and seam suppression against memory and throughput.

Either switch can be enabled alone. Enable both when physics and prior are both large enough to benefit; otherwise the collective overhead may outweigh the local work. Keep slurm_gres, slurm_ntasks_per_node, slurm_nodes, and the distribution switches coupled in a BenchOpt parameter grid. See the Configuration Guide for the coupled-list syntax.

Parameters#

Algorithm and prior#

denoiser (default: drunet): Learned denoiser used by the PnP prior. DRUNet is currently the only supported value; a 2D or 3D model is selected from the ground-truth shape.
denoiser_sigma (default: 0.05): Noise-level parameter \(\sigma\) passed to the denoiser at every iteration.
step_size (default: None): Gradient step size \(\gamma\). When omitted, the solver estimates the operator norm and uses step_size_scale / ||A||.
step_size_scale (default: 0.99): Safety factor applied only to the automatically computed step size.
denoiser_lambda_relaxation (default: None): If set to \(\lambda\), blends the gradient-step and denoised values with weight \(\alpha=\gamma\lambda/(1+\gamma\lambda)\). If omitted, the denoised value becomes the next iterate directly.
init_method (default: pseudo_inverse): Initial reconstruction. zeros creates a zero tensor; pseudo_inverse uses \(A^\dagger y\); and adjoint uses \(A^T y\). Operator-based initializations are peak-normalized and clipped when the dataset supplies a signal range.
norm_strategy (default: clip): clip clamps the denoiser output to the dataset range. dynamic maps that range to [0, 1] before denoising and maps it back afterward.

Distribution and memory#

distribute_physics (default: False): Distribute stacked physics and data-fidelity work across ranks.
distribute_denoiser (default: False): Distribute overlapping denoiser patches across ranks.
patch_size (default: 128): Patch extent along every tiled spatial dimension.
overlap (default: 32): Neighboring-patch overlap used for smooth reconstruction.
max_batch_size (default: 0): Maximum local patch batch passed to the denoiser. It is forwarded to the distributed processing layer; use a small positive value when memory-bound and increase it for throughput when memory permits.

Compilation#

compile (default: None): Placement of torch.compile(). pre compiles the denoiser and compatible forward/adjoint methods before distributed wrapping; post compiles the wrapped denoiser and data-fidelity gradient; fused compiles an entire gradient–denoise–clip iteration. fused requires norm_strategy: clip. The first compiled iteration includes compilation cost, so use profiler warmup iterations when measuring steady state.

Execution and reporting#

slurm_nodes (default: 1): Number of nodes requested for the solver run.
slurm_ntasks_per_node (default: 1): Number of distributed worker processes per node for SLURM execution.
slurm_gres (default: gpu:1): SLURM generic-resource request per node.
torchrun_nproc_per_node (default: 1): Processes per node when the run is launched through torchrun.
name_prefix (default: pnp): Prefix used to identify the run and profiler artifacts.
profiler_mode, profiler_warmup, profiler_active, profiler_trace_dir, profiler_per_step, profiler_repeat, profiler_save_file: Profiling backend, recording-window, and output options. See Profiling Solvers.