DistributedProcessing#

class deepinv.distributed.framework.DistributedProcessing(ctx, processor, *, strategy=None, strategy_kwargs=None, max_batch_size=None, checkpoint_batches='auto', checkpoint_use_reentrant=False, checkpoint_preserve_rng_state=True, **kwargs)[source]#

Bases: Module

Distributed signal processing using pluggable tiling and reduction strategies.

This class enables distributed processing of large signals (images, volumes, etc.) by:

  1. Splitting the signal into patches using a chosen strategy

  2. Distributing patches across multiple processes/GPUs

  3. Processing each patch independently using a provided processor function

  4. Combining processed patches back into the full signal with proper overlap handling

The processor can be any callable that operates on tensors (e.g., denoisers, priors, neural networks, etc.). The class handles all distributed coordination automatically.

Example:

import torch
from deepinv.distributed import DistributedContext
from deepinv.distributed.framework import DistributedProcessing

x = torch.randn(1, 3, 1024, 1024)

with DistributedContext() as ctx:
    processor = torch.nn.Identity()
    distributed_processor = DistributedProcessing(
        ctx,
        processor,
        strategy_kwargs={"patch_size": 256, "overlap": 32},
        max_batch_size=1,
    )
    output = distributed_processor(x.to(ctx.device))
Parameters:

  • ctx (DistributedContext) – distributed context manager.

  • processor (Callable[[torch.Tensor], torch.Tensor]) – processing function to apply to signal patches. Should accept a batched tensor of shape (B, C, ...) and return a tensor of the same shape. Examples: denoiser, neural network, prior gradient function, etc.

  • strategy (DistributedSignalStrategy | None) – signal processing strategy for patch extraction and reduction. Either a custom strategy instance or None, which corresponds to the default tiling strategy.

  • strategy_kwargs (dict | None) – additional keyword arguments passed to the strategy constructor when using string strategy names. Examples: patch_size, overlap, tiling_dims. Default is None.

  • max_batch_size (int | None) – maximum number of patches to process in a single batch. If None, all local patches are batched together. Set to 1 for sequential processing (useful for memory-constrained scenarios). Higher values increase throughput but require more memory. Default is None.

  • checkpoint_batches (str) – activation checkpointing mode for patch batches during the backward pass. Checkpointing saves memory by recomputing activations instead of storing them. Use "auto" (default) to enable it only when gradients are enabled and there are multiple local patch batches, "always" to enable it whenever gradients are enabled, or "never" to disable it.

  • checkpoint_use_reentrant (bool) – reentrant mode passed to torch.utils.checkpoint.checkpoint(). Default is False (recommended by PyTorch).

  • checkpoint_preserve_rng_state (bool) – whether to preserve RNG state across forward recomputation when checkpointing. Default is True.

forward(x, *args, gather=True, **kwargs)[source]#

Apply distributed processing to input signal.

Parameters:

  • x (torch.Tensor) – input signal tensor to process, typically of shape (B, C, H, W) for 2D or (B, C, D, H, W) for 3D signals.

  • args – additional positional arguments passed to the processor.

  • gather (bool) – whether to gather results across ranks. If False, returns local contribution. Default is True.

  • kwargs – additional keyword arguments passed to the processor.

Returns:

processed signal with the same shape as input.

Return type:

Tensor

Examples using DistributedProcessing:#

Distributed Denoiser with Image Tiling

Distributed Denoiser with Image Tiling

Distributed Plug-and-Play (PnP) Reconstruction

Distributed Plug-and-Play (PnP) Reconstruction