BSREM

class deepinv.optim.BSREM(data_fidelity=None, prior=None, lambda_reg=1.0, g_param=None, sigma_denoiser=None, num_subsets=2, stepsize=1.0, eps=1e-6, sensitivity_threshold=1e-2, max_iter=100, crit_conv='residual', thres_conv=1e-5, early_stop=False, custom_metrics=None, custom_init=None, unfold=False, trainable_params=None, cost_fn=None, params_algo=None, **kwargs)

Bases: BaseOptim

Block Sequential Regularized Expectation Maximization (BSREM) for Poisson inverse problems.

BSREM is a relaxed ordered-subsets algorithm for minimizing the penalized Poisson negative log-likelihood

\[\min_{x \in \mathbb{R}^n_{+}}\; \mathrm{KL}(y, Ax) + \lambda \reg{x}.\]

while preserving convergence guarantees [1][2]. With \(L\) subsets, one complete iteration of the algorithm applies the following update for \(l=1,\ldots,L\):

\[x_{k,l+1} = \mathcal{P}_{+}\left[x_{k,l} - \alpha_k \frac{x_{k,l}}{\bar{s}} \odot \left(\nabla f_l(x_{k,l}) + \frac{\lambda}{L}\nabla \reg{x_{k,l}}\right)\right],\]

where \(\bar{s}=A^T\mathbf{1}/L\) is the average subset sensitivity, \(\alpha_k\) is the relaxation step size, which is annealed over the iterations, and \(\mathcal{P}_{+}\) clamps the iterate to the positive orthant.

See deepinv.optim.optim_iterators.BSREMIteration for the details of one iteration.

Tip

The description of the algorithm above assumes unit Poisson gain. If a non-unit gain is used, the implementation automatically scales the preconditioner and prior to adapt to the gain.

A custom annealing schedule for the relaxation step size can be supplied as an iterable, for example stepsize=[1 / (1 + 0.1 * k) for k in range(max_iter)].

Note

The user can provide either the full measurement tensor y and full tomography physics, or pre-split measurements passed as a deepinv.utils.TensorList and pre-split physics passed as a deepinv.physics.StackedLinearPhysics. See deepinv.physics.split_physics() and deepinv.physics.split_measurements().

Note

By default, the algorithm is initialized with a tensor of ones with the same shape as \(A^T y\). This can be overridden using custom_init.

Parameters:

• num_subsets (int) – number of ordered subsets used for the splitting of the physics and measurements. Ignored when pre-split inputs are provided. Default: 2.

• data_fidelity (deepinv.optim.DataFidelity, list[DataFidelity]) – data fidelity used by the subset updates and to evaluate the objective. If None, defaults to deepinv.optim.PoissonLikelihood.

• prior (deepinv.optim.Prior, list[Prior]) – differentiable prior term. If None, no regularization is applied. Default: None.

• lambda_reg (float) – regularization parameter \(\lambda\). Default: 1.0.

• g_param (float) – parameter passed to the prior. Default: None.

• sigma_denoiser (float) – alias for g_param. If both are provided, g_param takes precedence. Default: None.

• stepsize (float, collections.abc.Iterable[float]) – annealing schedule of the relaxation step size. If an iterable is used, it must contain at least max_iter entries. Default: 1.0.

• eps (float) – positive value used for safe divisions and the positivity projection. Default: 1e-6.

• sensitivity_threshold (float) – relative sensitivity threshold defining the reconstruction support. Default: 1e-2.

• max_iter (int) – maximum number of BSREM epochs. Default: 100.

• crit_conv (str) – convergence criterion, either "residual" or "cost". Default: "residual".

• thres_conv (float) – convergence threshold for crit_conv. Default: 1e-5.

• early_stop (bool) – stop when the convergence criterion is met. Default: False.

• custom_metrics (dict) – custom metrics computed after every epoch. Default: None.

• custom_init (Callable) – custom initialization function. BSREM passes the split measurements and stacked subset physics to this function. If None, the reconstruction is initialized with ones. Default: None.

• unfold (bool) – whether to unfold the algorithm. Default: False.

• trainable_params (list) – algorithm parameters to train when unfolded, chosen from ["lambda", "stepsize", "g_param"]. If None, all parameters are trainable. Default: None.

• cost_fn (Callable) – custom cost function used for metrics and convergence. BSREM calls it with a deepinv.optim.StackedPhysicsDataFidelity, split measurements, and stacked subset physics. Default: None.

• params_algo (dict) – optional algorithm parameters. When provided, this overrides stepsize, lambda_reg, and g_param.

---

References:

[1]

A.R. De Pierro and M.E.B. Yamagishi. Fast EM-like methods for maximum “a posteriori” estimates in emission tomography. IEEE Transactions on Medical Imaging, 20(4):280–288, April 2001. doi:10.1109/42.921477.

[2]

Sangtae Ahn and Jeffrey A. Fessler. Globally convergent image reconstruction for emission tomography using relaxed ordered subsets algorithms. IEEE Transactions on Medical Imaging, 22(5):613–626, May 2003. doi:10.1109/TMI.2003.812251.

forward(y, physics, *args, **kwargs)

Run BSREM with full or pre-split measurements and physics.

Parameters:

• y (torch.Tensor, deepinv.utils.TensorList, list[torch.Tensor]) – Full measurement tensor, or pre-split measurements when physics is a deepinv.physics.StackedLinearPhysics.

• physics (deepinv.physics.Physics) – Full tomography/PET physics or pre-split deepinv.physics.StackedLinearPhysics.

Returns:

Reconstructed image, and optionally the metrics dictionary when compute_metrics=True.

Return type:

torch.Tensor | tuple[torch.Tensor, dict]

Examples using BSREM:

OSEM, BSREM and gradient descent for 2D BrainWeb PET

OSEM, BSREM and gradient descent for 2D BrainWeb PET

3D PET reconstruction with the Brainweb dataset

3D PET reconstruction with the Brainweb dataset