distribute#

deepinv.distributed.distribute(object, ctx, *, num_operators=None, type_object='auto', dtype=torch.float32, gather_strategy='concatenated', tiling_strategy=None, tiling_dims=None, patch_size=256, overlap=64, max_batch_size=None, checkpoint_batches='auto', checkpoint_use_reentrant=False, checkpoint_preserve_rng_state=True, **kwargs)[source]#

Distribute a DeepInverse object across multiple devices.

This function takes a DeepInverse object and distributes it using the provided DistributedContext.

The list of supported DeepInverse objects includes:

Parameters:

  • object (StackedPhysics | list[Physics] | Callable | Denoiser | DataFidelity | StackedPhysicsDataFidelity | list[DataFidelity] | torch.nn.Module | torch.nn.parameter.Parameter | Sequence[torch.nn.parameter.Parameter]) – DeepInverse object to distribute.

  • ctx (DistributedContext) – distributed context manager.

  • num_operators (int | None) – number of physics operators when using a factory for physics, otherwise inferred. Default is None.

  • type_object (str | None) – type of object to distribute. Options are 'physics', 'linear_physics', 'data_fidelity', 'denoiser', 'module', 'parameters', or 'auto' for automatic detection. Default is 'auto'. 'module' is restricted to BaseOptim models built with unfold=True (including the legacy BaseUnfold subclass). Generic torch.nn.Module instances are intentionally not supported by this API, to avoid ambiguous partial auto-distribution.

  • dtype (torch.dtype | None) – data type for distributed object. Default is torch.float32.

  • gather_strategy (str) –

    strategy for gathering distributed results.

    Options are:

    • 'naive': Simple object serialization (best for small tensors)

    • 'concatenated': Single concatenated tensor (best for medium/large tensors, minimal communication)

    • 'broadcast': Per-operator broadcasts (best for heterogeneous sizes or streaming)

    Default is 'concatenated'.

  • tiling_strategy (DistributedSignalStrategy | None) – strategy for tiling the signal (for Denoiser). Options are either a custom strategy instance or None, which corresponds to the default tiling strategy.

  • tiling_dims (int | tuple[int, ...] | None) –

    dimensions to tile over (for Denoiser).

    Can be one of the following:

    • If None (default), tiles the last N-2 dimensions of your input tensor.

    • If an int N, only tiles over the specified dimension.

    • If a tuple, specifies exact dimensions to tile.

    Examples:

    • For (B, C, H, W) image: tiling_dims=(2, 3) tiles over H and W.

    • For (B, C, D, H, W) volume: tiling_dims=(2, 3, 4) tiles over D, H, W.

    • For (B, C, H, W) image: tiling_dims=2 tiles only over H dimension.

    • For (B, C, D, H, W) volume: tiling_dims=None tiles over D, H, W dimensions.

  • patch_size (int) – size of patches for tiling strategies (for Denoiser). Can be an int (same size for all tiled dims) or a tuple (per-dimension size). Default is 256.

  • overlap (int) – receptive field size for overlap in tiling strategies (for Denoiser). Can be an int (same size for all tiled dims) or a tuple (per-dimension size). Default is 64.

  • max_batch_size (int | None) – maximum number of patches to process in a single batch (for Denoiser). If None, all patches are batched together. Set to 1 for sequential processing. Default is None.

  • checkpoint_batches (str) – activation checkpointing mode for patch-batches during backward (for Denoiser). Supported values are 'auto', 'always' and 'never'. Default is 'auto'.

  • checkpoint_use_reentrant (bool) – reentrant mode for activation checkpointing in denoiser processing. Default is False.

  • checkpoint_preserve_rng_state (bool) – preserve RNG state during checkpoint recomputation in denoiser processing. Default is True.

  • kwargs – additional keyword arguments for specific distributed classes.

Returns:

Distributed version of the input object.

Return type:

DistributedStackedPhysics | DistributedStackedLinearPhysics | DistributedProcessing | DistributedDataFidelity | Parameter | list[Parameter] | BaseOptim

Examples:

Distribute a Physics object:

>>> from deepinv.physics import Blur, StackedLinearPhysics
>>> from deepinv.distributed import DistributedContext, distribute
>>> with DistributedContext() as ctx:
...     physics = StackedLinearPhysics([Blur(kernel_size=5), Blur(kernel_size=9)])
...     dphysics = distribute(physics, ctx)

Distribute a DataFidelity object:

>>> from deepinv.optim.data_fidelity import L2
>>> from deepinv.distributed import DistributedContext, distribute
>>> with DistributedContext() as ctx:
...     data_fidelity = L2()
...     ddata_fidelity = distribute(data_fidelity, ctx)

Distribute a Prior object:

>>> from deepinv.models import DnCNN
>>> from deepinv.distributed import DistributedContext, distribute
>>> with DistributedContext() as ctx:
...     denoiser = DnCNN()
...     ddenoiser = distribute(denoiser, ctx)

Distribute a full unfolded PGD model in one call:

>>> from deepinv.models import DnCNN
>>> from deepinv.optim import PGD
>>> from deepinv.optim.data_fidelity import L2
>>> from deepinv.optim.prior import PnP
>>> from deepinv.distributed import DistributedContext, distribute
>>> with DistributedContext() as ctx:
...     model = PGD(
...         data_fidelity=L2(),
...         prior=PnP(DnCNN(in_channels=1, out_channels=1)),
...         stepsize=[0.9, 0.8],
...         max_iter=2,
...         unfold=True,
...     )
...     distribute(model, ctx, patch_size=64, overlap=8)

Examples using distribute:#

Distributed Denoiser with Image Tiling

Distributed Denoiser with Image Tiling

Distributed Physics Operators

Distributed Physics Operators

Distributed Plug-and-Play (PnP) Reconstruction

Distributed Plug-and-Play (PnP) Reconstruction

Distributed Training of Unfolded Networks

Distributed Training of Unfolded Networks