optim.FixedPoint

optim.FixedPoint#

class bartorch.optim.FixedPoint#

Bases: Module

Iteration block run to its fixed point and differentiated there (deep equilibrium).

The forward iterations run without recording until the state changes by at most tol of its norm, or for max_iter steps. The backward pass takes one recorded step from the fixed point and iterates the vector-Jacobian product of that step to its own fixed point, so memory does not grow with the iteration count.

Parameters:

  • block (nn.Module) – One of bartorch.optim’s iteration blocks.

  • max_iter (int, default=100) – Maximum number of forward iterations.

  • tol (float, default=1e-4) – Relative change of the state at which the forward iteration stops.

  • backward_iter (int, default=None) – Maximum number of backward iterations; max_iter when omitted.

  • backward_tol (float, default=None) – Relative change at which the backward iteration stops; tol when omitted.

iterations#

Forward iterations taken by the last call.

Type:

int

Raises:

  • TypeError – block is a FISTABlock, whose momentum makes every step a different map.

  • ValueError – block has no fixed point for another reason: a moving or scheduled rho, adaptive or relaxed step sizes, or an ImplicitPrior with a sigma schedule or step=True.

Examples

>>> deq = optim.FixedPoint(optim.PRIDUBlock(priors.ImplicitPrior(net, 0.05)), tol=1e-4)
>>> deq(kspace, A).abs().sub(target).square().sum().backward()
forward()#

Iterate the block from block.start(y, A, x0) and return the image at the fixed point.

The result is recorded for autograd only when y requires a gradient or the block has trainable parameters.

Parameters:

  • y (torch.Tensor) – Data.

  • A (LinearOperator) – Encoding operator.

  • x0 (torch.Tensor, default=None) – Starting image; zero when omitted.

Returns:

The image at the fixed point.

Return type:

torch.Tensor