Callback in Algorithm

[epapoutsellis]

During the Stochastic Project, I designed the following improvement for the callback used in the run method of our Algorithm class.

To instantiate an Algorithm we mainly pass instances of Functions and Operators. So I thought why not pass another class which will be responsible for the callbacks. At the moment, the callback used in the run method is restricted to iteration, x attributes and get_last_objective() method. https://github.com/TomographicImaging/CIL/blob/0418ed2abd42c085db25aa3d0405b4b16e6591e4/Wrappers/Python/cil/optimisation/algorithms/Algorithm.py#L286-L287 This is not flexible. It basically gives access to the user for only three attributes.

The design that I have is the following lines

https://github.com/epapoutsellis/StochasticCIL/blob/477d51fd94c9d6554ca143b86e4b95a21c60835d/Wrappers/Python/cil/optimisation/algorithms/Algorithm.py#L311-L313

Instead of using function that takes 3 inputs, I pass a Callable Class with acts on the algorithm which in this case is self, i.e., the instantiated algorithm. In this way, the user has access to every attribute/method of the Algorithm class.

What the users can do with it? Anything they want

How I use it for the Stochastic Project?

• I created a class that computes some Quality Metrics. It is basically a class that computes PSNR, MSE, MAE given a reference image, e.g., Metric(current iterate, reference). Users can pass other metrics from other libraries, e.g., SSIM, NRMSE from scikit-image without reimplementing everything. In addition, an attribute is created on the fly, so the user can have access to these values, e.g., ista.ssim will be a list of all the ssim values computed every update_objective_interval. Finally, verbose allows them to print or not all this information.

• Instead of computing metrics and statistics of all the pixels/voxels, our users can define a specific ROI. This is useful for PET/SPECT applications measuring Standardized Uptake Value (SUV). See here @gschramm . For instance we can compute the SSIM for all the pixels/voxels and the .mean for some ROI at the same time.

• In some cases where our objectives are smooth, it is useful to compute the $|\nabla f (x{k})||$ instead of the objective $f(x{k})$

class ComputeGradient(AlgorithmDiagnostics):

    def __init__(self):
        super(ComputeGradient, self).__init__(verbose = 0)        

    def __call__(self, algo):
        if not hasattr(algo, "gradf"):
            setattr(algo, "gradf", [])                
        algo.gradf.append(algo.f.gradient(algo.get_output()).norm())

and the user can do plt.plot(ista.gradf) to get

• Also, we can compute other distances $||x{k} - x^{*}||{2}$.

• In the Algorithm class the default stopping criterion is max_iteration. We can replace it using a class that overrides the should_stop method. For instance, we can use $||x{k+1} - x{k}||_{2}<\mbox{tol}$. Also, in some cases, e.g., FISTA/PDHG with an iterative inner solver (FGP_TV, TotalVariation, TGV) the important is not only on the cold/warm start but also on the metric and accuracy used to solve the inner problem. For instance, the following rule https://github.com/TomographicImaging/CIL/blob/0418ed2abd42c085db25aa3d0405b4b16e6591e4/Wrappers/Python/cil/optimisation/functions/TotalVariation.py#L333-L338 is kind of strange in my opinion, but I implemented to have an exact match with the FGP_TV from CCPi-Regtk. According to this paper there are better stopping rules.

• We can also define stopping rules using algo.timings and compared different algorithms in terms of wall time, e.g., Difference from optimal vs computational time. Particularly useful for benchmarking stochastic vs deterministic recons.

• One can create fancy dynamic plots, for instance a row of 3 figures showing the current iterate, difference from optimal solution and an SSIM plot (iterations vs SSIM values).

• We can use Weights and Biases platform to monitor whatever we want when an algorithm is executed.

  

• There is no need to have a logger in the Algorithm Class

https://github.com/TomographicImaging/CIL/blob/0418ed2abd42c085db25aa3d0405b4b16e6591e4/Wrappers/Python/cil/optimisation/algorithms/Algorithm.py#L83-L90

A Logger class passed to callback can take care of this. Also, wandb has this functionality. @gschramm used Tensorboard here. I had some problems with Tensorboard and its gui.

Finally, a similar design is used for the Proximal Gradient Algorithm, Preconditioner and StepSize classes, but I will create separate issues for them.

Note Atm, all the callbacks are executed every update_objective_interval. This can be easily fixed if we want but is not important in my opinion. Basically, we can compute SSIM every 10 iterations, plot the current iterate every 2 iterations and the .mean of $x_{k}$ every 50 iterations.

[casperdcl]

I guess maybe we should have a new callbacks: Iterable[callable] kwarg (leaving the existing callback: callable kwarg as-is for backwards compatibility)? /CC @paskino

[casperdcl]

Had a quick chat with @paskino @MargaretDuff.

I'm going to follow https://keras.io/api/callbacks design and:

• introduce a new base Callback class
• put existing logging functionality into new callback classes inheriting from Callback
• add new functionality too (e.g. tqdm-based progress bar callbacks)
• introduce a new callbacks=[] kwarg
• document how to create/customise callbacks
• soft-deprecate the old callback/verbose kwargs
  • if not None, print a warning and construct a new Callback class based on them

[epapoutsellis]

Are you going to move update_objective_interval from the signature of Algorithm and pass it to Callback?

[paskino]

Discussed with @epapoutsellis @MargaretDuff @jakobsj

It seems that the callback route could be used to define the stopping criterion?

[MargaretDuff]

Discussed with @epapoutsellis @MargaretDuff @jakobsj

It seems that the callback route could be used to define the stopping criterion?

Callbacks acting on the algorithm class can be used to raise StopIteration() and then terminate the iterations. Keras has an EarlyStopping class which stops training when the monitored metric stops improving https://keras.io/api/callbacks/early_stopping/.

There was a bit of a discussion with two options:

• We encourage the callbacks argument to be used for just observing an algorithm, and not interfering and have another argument for stopping criterion. This would be very similar to the callbacks argument in that it would take a class that when called acts on the algorithm class.
• The user will be able to pass a callback which stops the iterations whether we want them to or not. With good documentation and examples, we should encourage everything to be passed to callbacks.