Source code for deepinv.optim.fixed_point

import torch
import torch.nn as nn
import warnings
from tqdm import tqdm


[docs] class FixedPoint(nn.Module): """ Fixed-point iterations module. This module implements the fixed-point iteration algorithm given a specific fixed-point iterator (e.g. proximal gradient iteration, the ADMM iteration, see :meth:`deepinv.optim.optim_iterators`), that is for :math:`k=1,2,...` .. math:: \qquad (x_{k+1}, u_{k+1}) = \operatorname{FixedPoint}(x_k, u_k, f, g, A, y, ...) \hspace{2cm} (1) where :math:`f` is the data-fidelity term, :math:`g` is the prior, :math:`A` is the physics model, :math:`y` is the data. :Examples: This example shows how to use the :class:`FixedPoint` class to solve the problem :math:`\min_x 0.5*||Ax-y||_2^2 + \lambda*||x||_1` with the PGD algorithm, where A is the identity operator, :math:`\lambda = 1` and :math:`y = [2, 2]`. >>> import deepinv as dinv >>> # Create the measurement operator A >>> A = torch.tensor([[1, 0], [0, 1]], dtype=torch.float64) >>> A_forward = lambda v: A @ v >>> A_adjoint = lambda v: A.transpose(0, 1) @ v >>> # Define the physics model associated to this operator >>> physics = dinv.physics.LinearPhysics(A=A_forward, A_adjoint=A_adjoint) >>> # Define the measurement y >>> y = torch.tensor([2, 2], dtype=torch.float64) >>> # Define the data fidelity term >>> data_fidelity = dinv.optim.data_fidelity.L2() >>> # Define the prior term >>> prior = dinv.optim.prior.L1Prior() >>> # Define the parameters of the algorithm >>> params_algo = {"g_param": 1.0, "stepsize": 1.0, "lambda": 1.0, "beta": 1.0} >>> # Choose the iterator associated to the PGD algorithm >>> iterator = dinv.optim.optim_iterators.PGDIteration() >>> # Iterate the iterator >>> x_init = torch.tensor([2, 2], dtype=torch.float64) # Define initialisation of the algorithm >>> X = {"est": (x_init ,), "cost": []} # Iterates are stored in a dictionary of the form {'est': (x,z), 'cost': F} >>> max_iter = 50 >>> for it in range(max_iter): ... X = iterator(X, data_fidelity, prior, params_algo, y, physics) >>> # Return the solution >>> X["est"][0] tensor([1., 1.], dtype=torch.float64) :param deepinv.optim.optim_iterators.optim_iterator iterator: function that takes as input the current iterate, as well as parameters of the optimization problem (prior, measurements, etc.) :param function update_params_fn: function that returns the parameters to be used at each iteration. Default: ``None``. :param function update_prior_fn: function that returns the prior to be used at each iteration. Default: ``None``. :param function init_iterate_fn: function that returns the initial iterate. Default: ``None``. :param function init_metrics_fn: function that returns the initial metrics. Default: ``None``. :param function check_iteration_fn: function that performs a check on the last iteration and returns a bool indicating if we can proceed to next iteration. Default: ``None``. :param function check_conv_fn: function that checks the convergence after each iteration, returns a bool indicating if convergence has been reached. Default: ``None``. :param int max_iter: maximum number of iterations. Default: ``50``. :param bool early_stop: if True, the algorithm stops when the convergence criterion is reached. Default: ``True``. :param bool anderson_acceleration: if True, the Anderson acceleration is used. Default: ``False``. :param int history_size: size of the history used for the Anderson acceleration. Default: ``5``. :param float beta_anderson_acc: momentum of the Anderson acceleration step. Default: ``1.0``. :param float eps_anderson_acc: regularization parameter of the Anderson acceleration step. Default: ``1e-4``. """ def __init__( self, iterator=None, update_params_fn=None, update_data_fidelity_fn=None, update_prior_fn=None, init_iterate_fn=None, init_metrics_fn=None, update_metrics_fn=None, check_iteration_fn=None, check_conv_fn=None, max_iter=50, early_stop=True, anderson_acceleration=False, history_size=5, beta_anderson_acc=1.0, eps_anderson_acc=1e-4, verbose=False, show_progress_bar=False, ): super().__init__() self.iterator = iterator self.max_iter = max_iter self.early_stop = early_stop self.update_params_fn = update_params_fn self.update_data_fidelity_fn = update_data_fidelity_fn self.update_prior_fn = update_prior_fn self.init_iterate_fn = init_iterate_fn self.init_metrics_fn = init_metrics_fn self.update_metrics_fn = update_metrics_fn self.check_conv_fn = check_conv_fn self.check_iteration_fn = check_iteration_fn self.anderson_acceleration = anderson_acceleration self.history_size = history_size self.beta_anderson_acc = beta_anderson_acc self.eps_anderson_acc = eps_anderson_acc self.verbose = verbose self.show_progress_bar = show_progress_bar if self.check_conv_fn is None and self.early_stop: warnings.warn( "early_stop is set to True but no check_conv_fn has been defined." ) self.early_stop = False
[docs] def init_anderson_acceleration(self, X): r""" Initialize the Anderson acceleration algorithm. Code inspired from `this tutorial <http://implicit-layers-tutorial.org/deep_equilibrium_models/>`_. :param dict X: initial iterate. """ x = X["est"][0] b, d, h, w = x.shape x_hist = torch.zeros( b, self.history_size, d * h * w, dtype=x.dtype, device=x.device ) # history of iterates. T_hist = torch.zeros( b, self.history_size, d * h * w, dtype=x.dtype, device=x.device ) # history of T(x_k) with T the fixed point operator. H = torch.zeros( b, self.history_size + 1, self.history_size + 1, dtype=x.dtype, device=x.device, ) # H in the Anderson acceleration linear system Hp = q . H[:, 0, 1:] = H[:, 1:, 0] = 1.0 q = torch.zeros( b, self.history_size + 1, 1, dtype=x.dtype, device=x.device ) # q in the Anderson acceleration linear system Hp = q . q[:, 0] = 1 return x_hist, T_hist, H, q
[docs] def anderson_acceleration_step( self, it, X_prev, TX_prev, x_hist, T_hist, H, q, cur_data_fidelity, cur_prior, cur_params, *args, ): r""" Anderson acceleration step. Code inspired from `this tutorial <http://implicit-layers-tutorial.org/deep_equilibrium_models/>`_. :param int it: current iteration. :param dict X_prev: previous iterate. :param dict TX_prev: output of the fixed-point operator evaluated at X_prev :param torch.Tensor x_hist: history of last ``history-size`` iterates. :param torch.Tensor T_hist: history of T evlauaton at the last ``history-size``, where T is the fixed-point operator. :param torch.Tensor H: H in the Anderson acceleration linear system Hp = q . :param torch.Tensor q: q in the Anderson acceleration linear system Hp = q . :param deepinv.optim.DataFidelity cur_data_fidelity: Instance of the DataFidelity class defining the current data_fidelity. :param deepinv.optim.prior cur_prior: Instance of the Prior class defining the current prior. :param dict cur_params: Dictionary containing the current parameters of the algorithm. :param args: arguments for the iterator. """ x_prev = X_prev["est"][0] # current iterate Tx Tx_prev = TX_prev["est"][0] # current iterate x b = x_prev.shape[0] # batchsize x_hist[:, it % self.history_size] = x_prev.reshape( b, -1 ) # prepare history of x T_hist[:, it % self.history_size] = Tx_prev.reshape( b, -1 ) # prepare history of Tx m = min(it + 1, self.history_size) G = T_hist[:, :m] - x_hist[:, :m] H[:, 1 : m + 1, 1 : m + 1] = ( torch.bmm(G, G.transpose(1, 2)) + self.eps_anderson_acc * torch.eye(m, dtype=Tx_prev.dtype, device=Tx_prev.device)[None] ) p = torch.linalg.solve(H[:, : m + 1, : m + 1], q[:, : m + 1])[ :, 1 : m + 1, 0 ] # solve the linear system H p = q. x = ( self.beta_anderson_acc * (p[:, None] @ T_hist[:, :m])[:, 0] + (1 - self.beta_anderson_acc) * (p[:, None] @ x_hist[:, :m])[:, 0] ) # Anderson acceleration step. x = x.view(x_prev.shape) F = ( self.iterator.F_fn(x, cur_data_fidelity, cur_prior, cur_params, *args) if self.iterator.has_cost else None ) est = list(TX_prev["est"]) est[0] = x return {"est": est, "cost": F}
[docs] def forward(self, *args, compute_metrics=False, x_gt=None, **kwargs): r""" Loops over the fixed-point iterator as (1) and returns the fixed point. The iterates are stored in a dictionary of the form ``X = {'est': (x_k, u_k), 'cost': F_k}`` where: * ``est`` is a tuple containing the current primal and auxiliary iterates, * ``cost`` is the value of the cost function at the current iterate. Since the prior and parameters (stepsize, regularisation parameter, etc.) can change at each iteration, the prior and parameters are updated before each call to the iterator. :param bool compute_metrics: if ``True``, the metrics are computed along the iterations. Default: ``False``. :param torch.Tensor x_gt: ground truth solution. Default: ``None``. :param args: optional arguments for the iterator. Commonly (y,physics) where ``y`` (torch.Tensor y) is the measurement and ``physics`` (deepinv.physics) is the physics model. :param kwargs: optional keyword arguments for the iterator. :return tuple: ``(x,metrics)`` with ``x`` the fixed-point solution (dict) and ``metrics`` the computed along the iterations if ``compute_metrics`` is ``True`` or ``None`` otherwise. """ X = ( self.init_iterate_fn(*args, F_fn=self.iterator.F_fn) if self.init_iterate_fn else None ) metrics = ( self.init_metrics_fn(X, x_gt=x_gt) if self.init_metrics_fn and compute_metrics else None ) self.check_iteration = True if self.anderson_acceleration: self.x_hist, self.T_hist, self.H, self.q = self.init_anderson_acceleration( X ) it = 0 for it in tqdm( range(self.max_iter), disable=(not self.verbose or not self.show_progress_bar), ): X_prev = X X = self.single_iteration( X, it, *args, **kwargs, ) if self.check_iteration: metrics = ( self.update_metrics_fn(metrics, X_prev, X, x_gt=x_gt) if self.update_metrics_fn and compute_metrics else None ) if ( self.early_stop and (self.check_conv_fn is not None) and it > 1 and self.check_conv_fn(it, X_prev, X) ): break it += 1 return X, metrics
def single_iteration(self, X, it, *args, **kwargs): cur_params = self.update_params_fn(it) if self.update_params_fn else None cur_data_fidelity = ( self.update_data_fidelity_fn(it) if self.update_data_fidelity_fn else None ) cur_prior = self.update_prior_fn(it) if self.update_prior_fn else None X_prev = X X = self.iterator( X_prev, cur_data_fidelity, cur_prior, cur_params, *args, **kwargs ) if self.anderson_acceleration: X = self.anderson_acceleration_step( it, X_prev, X, self.x_hist, self.T_hist, self.H, self.q, cur_data_fidelity, cur_prior, cur_params, *args, ) self.check_iteration = ( self.check_iteration_fn(X_prev, X) if self.check_iteration_fn else True ) return X if self.check_iteration else X_prev