Fixes #78 #7574 #10973
Continuation of the work started in #16925 by @rth.

Goal

The goal of this PR is to propose a callback API that can handle the most important / asked use cases.

• Convergence monitoring (IterativeImputer not converging (at all) #14338 DictionaryLearning should support more general constraints on the components  #8994 (comment))
  Monitor some quantities / metrics at each iteration.
  This can also be very useful for maintenance / debugging / implementation of new features.
• Early stopping (Feature discussion: callbacks for long operations #10973)
  Allow to stop the iterations based on some external metric evaluated on a validation set.
• Progress bars ([Feature Request] Progress Bars #7574 Use python logging to report on convergence progress it level info for long running tasks #78 Feature discussion: callbacks for long operations #10973)
• Logging / Verbose (Use python logging to report on convergence progress it level info for long running tasks #78)
• Snapshots
  Take regular snapshots of an estimator during the fit to be able to recover a working estimator if the fit is somehow interrupted for instance.
• Resource monitoring
  (not implemented yet)

Challenges

• Supporting all these features and make each of these callbacks available is not easy and will require some refactoring in probably many estimators.

  The proposed API makes it possible to enable the callbacks 1 estimator at a time: Setting callbacks on non yet supported estimators has no effect. Thus we can then incrementally do it in subsequent dedicated PRs. Here I only did NMF, LogisticRegression and Pipeline to show what are the necessary changes in the code base.

  The proposed API also makes it possible to only enable a subset of the features for an estimator, and add the remaining ones later. For LogisticRegression I only passed the minimum for instance.

• Callbacks should not impact the performance of the estimators. Some quantities passed to the callbacks might be costly to compute. We don't want to spend time computing them if the only callback is a progress bar for instance.

  The solution I found is to do a lazy evaluation using lambdas and only actually compute them if there's at least 1 callback requesting it. For now callbacks can request these by defining specific class attributes but maybe there's a better way. mixins ?

• The callbacks described above are not meant to be evaluated a the same fitting step of an estimator.

  When an estimator has several nested loops (LogisticRegressionCV(multiclass="ovr") for instance has a loop over Cs, a loop over the classes and then the final loop for the iterations on the dataset), the snapshot callback can only be evaluated at the end of an outermost loop while the EarlyStopping would be evaluated at the end of an innermost loop, and the ProgressBar could be evaluated at each level of nesting.

  In this PR I propose that each estimator holds a computation tree as a private attribute representing these nested loops, the root being the beginning of fit and each node being one step of a loop. This structure is defined in _computation_tree.py. It allows to have a simple way to know exactly at which step of the fit we are at each evaluation of the callbacks and is kind of the best solution I found to solve the challenges described below. This imposes the main changes to the code base, i.e. passing the parent node around.

• Dealing with parallelism and especially multiprocessing is the main challenge to me.

  Typically with a callback you might want to accumulate a bunch of info during fit and recover them at the end. The issue is that the callback is not shared between sub-processes and modifying its state in a sub-process (e.g. modifying an attribute) will not be visible from the main process. The joblib API doesn't allow inter-process communication that would be needed to overcome this.

  The solution we found is that the callbacks write the information they want to keep in files (in files in this first implementation but we might consider sockets or another solution ?). It's relatively easy to avoid race conditions with this design.
  As an example this is necessary to be able to report progress in real time. In an estimator running in parallel, there's no like current computation node. We are at different nodes at the same time. But having the status of each node in a file updated at each call to the callbacks allows to know the current overall progress from the main process. (there are other difficulties described later).

• The last main challenge is meta-estimators. We'd like some callbacks to be set on the meta-estimator, like progress bars, but some others to be set on the underlying estimator(s), like early-stopping. Moreover we encounter the parallelism issue again if the meta-estimator is fitting clones of the underlying estimator in parallel, like GS.

  For that, I propose to have a mixin to tell a callback that it should be propagated to sub estimators. This way the meta-estimator will only propagate the appropriate callbacks to its sub-estimators, and these sub-estimators can also have normal callbacks.

The API

This PR adds a new module sklearn.callback which exposes BaseCallback, the abstract base class for the callbacks. All callbacks must inherit from BaseCallback. It also exposes AutoPropagatedMixin. Callbacks that should be propagated to sub-estimators by meta-estimators must inherit from this.

BaseCallback has 3 abstract methods:

• on_fit_end. Called at the beginning of fit, after all validations. We pass a reference to the estimator, X_train and y_train.
• on_fit_iter_end. Called at the end of each node of the computation tree, i.e. each step of each nested loop. We pass a reference to the estimator (which at this point might be different from the one passed at in_fit_begin for propagated callbacks), and the computation node where it was called. We also pass some of these:
  • stopping_criterion: when the estimator has a stopping criterion such that the iterations stop when stopping_criterion <= tol.
  • tol: tolerance for the stopping criterion.
  • reconstruction_attributes. These are the necessary attributes to construct an estimator (by copying the estimator and setting these as attributes) which will behave as if the fit stopped at this node. Then we must be able to call predict, transform, ...
  • fit_state: Model specific quantities updated during fit. This is not meant to be used by generic callbacks but by a callback designed for a specific estimator instead. This arg is not used in any of the use cases described above but I thinkit's important to have for custom callbacks. It's the role of each estimator to decide what is interesting to pass to the callback. We could later think of new field in the docstring of the estimators to describe what keys they pass to this arg.
• on_fit_end. Called at the end of fit. Takes no argument. It allows the callback to do some clean-up.

Examples

• Progress bars.

  expand

  Here's an example of progress monitoring using rich. I used custom estimators to simulate a complex setting with a meta-estimator (like a GridSearchCV) running in parallel with a sub-estimator also running in parallel.

  simplescreenrecorder-2021-12-16_19.17.35.mp4

• Convergence Monitoring

  expand

  from sklearn.decomposition import NMF
  import numpy as np
  X = np.random.random_sample((1100, 100))
  X_val = X[-100:]
  nmf = NMF(n_components=20, solver="mu")
  callback = ConvergenceMonitor(X_val=X_val)
  nmf._set_callbacks(callback)
  nmf.fit(X[:1000])
  callback.plot()

  

• Snapshot

  expand

  from sklearn.decomposition import NMF
  import numpy as np
  X = np.random.random_sample((1100, 100))
  nmf = NMF(n_components=20, solver="mu")
  callback = Snapshot()
  nmf._set_callbacks(callback)
  nmf.fit(X[:1000])
  # interrupt fit. Ctrl-C for instance
  # [...]
  KeyboardInterrupt:
  
  import pickle
  with open(callback.directory / "2021-12-16_19-33-15-083014.pkl", "rb") as f:
      new_nmf = pickle.load(f)
  W = new_nmf.transform(X[-100:])

• EarlyStopping

  expand

  If the on_fit_iter_end method of the callbacks returns True, the iteration loop breaks.

  from sklearn.decomposition import NMF
  import numpy as np
  X = np.random.random_sample((1100, 100))
  X_val = X[-100:]
  nmf = NMF(n_components=20, solver="mu")
  callback = EarlyStopping(monitor="objective_function", X_val=X_val, max_no_improvement=10, tol=1e-4)
  nmf._set_callbacks(callback)
  nmf.fit(X[:1000])

• Verbose

  expand

  from sklearn.decomposition import NMF
  import numpy as np
  X = np.random.random_sample((1100, 100))
  nmf = NMF(n_components=20, solver="mu", max_iter=20)
  nmf._set_callbacks(TextVerbose())
  nmf.fit(X)

  [NMF] iter 0 | time 0.02493s | stopping_criterion=8.730E-01 | tol=1.000E-04
  [NMF] iter 1 | time 0.02634s | stopping_criterion=8.737E-01 | tol=1.000E-04
  [NMF] iter 2 | time 0.02768s | stopping_criterion=8.743E-01 | tol=1.000E-04
  [NMF] iter 3 | time 0.02893s | stopping_criterion=8.749E-01 | tol=1.000E-04
  [NMF] iter 4 | time 0.03016s | stopping_criterion=8.755E-01 | tol=1.000E-04
  [NMF] iter 5 | time 0.03136s | stopping_criterion=8.760E-01 | tol=1.000E-04
  [NMF] iter 6 | time 0.03255s | stopping_criterion=8.766E-01 | tol=1.000E-04
  [NMF] iter 7 | time 0.03375s | stopping_criterion=8.772E-01 | tol=1.000E-04
  [NMF] iter 8 | time 0.03496s | stopping_criterion=8.777E-01 | tol=1.000E-04
  [NMF] iter 9 | time 0.03691s | stopping_criterion=8.782E-01 | tol=1.000E-04
  [NMF] iter 10 | time 0.03841s | stopping_criterion=5.307E-04 | tol=1.000E-04
  [NMF] iter 11 | time 0.03966s | stopping_criterion=1.049E-03 | tol=1.000E-04
  [NMF] iter 12 | time 0.04087s | stopping_criterion=1.552E-03 | tol=1.000E-04
  [NMF] iter 13 | time 0.04209s | stopping_criterion=2.036E-03 | tol=1.000E-04
  [NMF] iter 14 | time 0.04327s | stopping_criterion=2.498E-03 | tol=1.000E-04
  [NMF] iter 15 | time 0.04447s | stopping_criterion=2.936E-03 | tol=1.000E-04
  [NMF] iter 16 | time 0.04565s | stopping_criterion=3.349E-03 | tol=1.000E-04
  [NMF] iter 17 | time 0.04686s | stopping_criterion=3.734E-03 | tol=1.000E-04
  [NMF] iter 18 | time 0.04804s | stopping_criterion=4.093E-03 | tol=1.000E-04
  [NMF] iter 19 | time 0.04923s | stopping_criterion=4.425E-03 | tol=1.000E-04 
  

TODO

This PR is still WIP.

• It's missing all the documentation of the callback module to describe the API and how to use and write callbacks and an example.
• I started adding tests for the computation tree but we need more, and I still need to add test for the callback api and tests for each of the implemented callbacks.
• Finalize and Document the implemented callbacks. There are still a few issues that need to be fixed in these callbacks.
• Think about how callbacks should be reinitialized when reused, like refitting an estimator.