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:

2D PET reconstruction with the Brainweb dataset

2D PET reconstruction with the Brainweb dataset

3D PET reconstruction with the Brainweb dataset

3D PET reconstruction with the Brainweb dataset