priors.ImplicitPrior#

class bartorch.priors.ImplicitPrior#

Bases: Module

Denoiser in place of a regularization term’s proximal operator (plug-and-play).

The proximal step is denoiser(x, sigma), or denoiser(x) without a sigma, independently of the solver’s step size. Unlike BART’s proximal operators it is differentiable, so a solve containing it can be differentiated end to end. The proximal solvers of bartorch.optim take it in place of a term.

With transform, the term is \(g(G x)\): the denoiser is applied on the codomain of \(G\), and the alternating-direction and primal-dual iterations split the variable at \(G x\), as for a BART term with a transform. This places a prior learned in another representation, such as contrast-weighted images synthesized from subspace coefficients.

Parameters:

denoiser (callable) – Called as denoiser(x) or denoiser(x, sigma), with step=k added under step. Without spatial, x is the complex tensor (batch, *shape), with shape the image’s or the codomain of \(G\); with spatial, a network on real (n, channels, *spatial) tensors of order unity, such as a deepinv or monai model.
sigma (float or sequence of float, default=None) – Noise level passed to the denoiser, in its own units. A sequence is a schedule, one per iteration, the last repeated beyond its end. Held as a frozen torch.nn.Parameter.
transform (LinearOperator, default=None) – \(G\), with domain the image; the identity when omitted.
spatial ({2, 3}, default=None) – Trailing axes the network takes; the axes in front are folded into its batch axis. None passes the complex tensor unconverted.
channels ({1, 2, 3}, default=1) – Input channels of the network: 1 for grayscale, 3 for RGB (the plane is repeated and the three outputs averaged), 2 for the real and imaginary parts jointly. Only with spatial.
parts ({"channels", "separate", "magnitude"}, default=None) – Real and imaginary parts as two channels, each part denoised separately in one call, or the modulus denoised with the phase kept. "channels" for two channels and "separate" otherwise when omitted. Only with spatial.
normalize (bool, default=True) – Divide each image by its peak modulus around the call, so sigma is relative to that peak. Only with spatial.
step (bool, default=False) – Pass the iteration index to the denoiser as step.

Raises:

TypeError – denoiser is not callable.
ValueError – channels, parts or normalize without spatial, or an inconsistent combination of them; a sigma that is empty or not one-dimensional.

Notes

A sigma schedule or step=True makes each iteration a different map (annealed plug-and-play, or a network shared by the iterations of an unrolled reconstruction, such as bartorch.learning.UNet with steps=True); bartorch.optim.FixedPoint refuses both.

Examples

>>> from deepinv.models import DRUNet
>>> prior = priors.ImplicitPrior(DRUNet(in_channels=1, out_channels=1), sigma=0.05, spatial=2)
>>> image = optim.fista(y, A, prior)
>>> priors.ImplicitPrior(denoiser, transform=linop.MultiplySum(basis, ...))

property stationary#

no schedule and no index.

Type:: Whether every iteration applies the same denoiser

prox()#

Apply the denoiser to x; gamma is ignored.

Parameters:

x (tensor) – Of prox_shape(), or with a leading batch axis.
gamma (float, default=1.0) – Ignored: the denoising strength is sigma.
image_shape (tuple of int, default=None) – Image shape, which tells a single x from a batch; without it x is a single item.
iteration (int, default=0) – Iteration index, which selects the sigma of a schedule and is passed as step under step=True.

Returns:

Of x’s shape.

Return type:

torch.Tensor

prox_shape()#: Shape the denoiser is applied on: the image’s, or the codomain of transform.

apply_transform()#

Apply \(G\), \(G^H\) or \(G^H G\) to x, recorded for autograd.

The identity without transform.

Parameters:

x (torch.Tensor) – Input of the mode, with or without a leading batch axis, applied item by item.
image_shape (tuple of int, default=None) – Image shape, checked against the domain of \(G\).
mode ({"forward", "adjoint", "normal"}, default="forward") – \(G\), \(G^H\) or \(G^H G\).

Return type: