{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Supporting callbacks in third party estimators\n\n.. currentmodule:: sklearn.callback\n\nThis document shows how to make third party :term:`estimators` and\n:term:`meta-estimators` compatible with the callback infrastructure supported by\nscikit-learn.\n\nGenerally speaking, a callback is a function that is provided by the user to be called\nat specific steps of a process, or to be triggered by specific events. Callbacks provide\na clean mechanism for inserting custom logic like monitoring progress or metrics,\nwithout modifying the core algorithm of the process.\n\nIn scikit-learn, callbacks take the form of classes following a [protocol](https://typing.python.org/en/latest/spec/protocol.html)_. This protocol requires the\ncallback classes to implement specific methods (referred to as callback hooks) which\nare called at specific steps of the fitting of an estimator or a meta-estimator.\nThese hooks are :meth:`~FitCallback.setup`, :meth:`~FitCallback.on_fit_task_begin`,\n:meth:`~FitCallback.on_fit_task_end` and :meth:`~FitCallback.teardown`. The\n:meth:`~FitCallback.setup` and :meth:`~FitCallback.teardown` hooks are called only once,\nrespectively at the start and end of the estimator's :term:`fit` method, and are\nresponsible for setting up and shutting down the callback. The\n:meth:`~FitCallback.on_fit_task_begin` and :meth:`~FitCallback.on_fit_task_end` hooks\nare respectively called at the beginning and end of each task in `fit` and are\nresponsible for the actual callback work. In scikit-learn estimators, a task in `fit` is\nusually one step of a loop, with nested loops corresponding to nested tasks. In general,\na task can be whatever unit of work the estimator's developer wants it to be.\n\nIn order to support the callbacks, estimators need to initialize and manage\n:class:`~CallbackContext` objects. As the name implies, these objects hold the\ncontextual information necessary to run the callback hooks. They are also responsible\nfor calling the callback hooks at the right time.\n\nIn the following, we show how to convert an example estimator class and an example\nmeta-estimator class to make them compliant with the scikit-learn callback\ninfrastructure.\n\nFirst a few imports and some random data for the rest of the script.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Authors: The scikit-learn developers\n# SPDX-License-Identifier: BSD-3-Clause" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import numpy as np\n\nfrom sklearn.base import BaseEstimator, clone\nfrom sklearn.callback import CallbackSupportMixin, ProgressBar, with_callbacks\nfrom sklearn.metrics.pairwise import euclidean_distances\nfrom sklearn.model_selection import check_cv\nfrom sklearn.utils import check_random_state\nfrom sklearn.utils.parallel import Parallel, delayed\nfrom sklearn.utils.validation import check_is_fitted, validate_data\n\nn_samples, n_features = 100, 4\nrng = np.random.RandomState(42)\nX = rng.rand(n_samples, n_features)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Custom Estimator\nHere we demonstrate how to implement a custom estimator that supports callbacks. For\nthe example, a simplified version of KMeans is presented. First, let's implement our\n`SimpleKMeans` estimator without the callback support.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "class SimpleKMeans(BaseEstimator):\n def __init__(self, n_clusters=6, n_iter=100, random_state=None):\n self.n_clusters = n_clusters\n self.n_iter = n_iter\n self.random_state = random_state # controls the centroids' initialization\n\n def _compute_labels(self, X):\n # Get the index of the closest centroid for each point in X.\n return np.argmin(euclidean_distances(X, self.cluster_centers_), axis=1)\n\n def fit(self, X, y=None):\n # `y` is not used but we need to declare it to adhere to scikit-learn's\n # estimators fit convention.\n\n # Input validation is a good practice in estimators, for more information about\n # it you can refer to\n # https://scikit-learn.org/stable/developers/develop.html#input-validation.\n X = validate_data(self, X)\n random_state = check_random_state(self.random_state)\n # Randomly initialize the centroids.\n self.cluster_centers_ = random_state.rand(self.n_clusters, X.shape[1])\n\n for i in range(self.n_iter):\n # The fit iterations consist in getting the cluster label of each data point\n # according to their closest centroid, and then updating the centroids as\n # the center of each cluster.\n labels = self._compute_labels(X)\n\n for k in range(self.n_clusters):\n # For each centroid, if its cluster is not empty, its coordinates are\n # updated with the coordinates of the cluster's center.\n if (labels == k).any():\n self.cluster_centers_[k] = X[labels == k].mean(axis=0)\n\n return self\n\n def predict(self, X):\n check_is_fitted(self)\n return self._compute_labels(X)\n\n def transform(self, X):\n check_is_fitted(self)\n return euclidean_distances(X, self.cluster_centers_)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now let's add all the elements necessary to support callbacks.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# First things first, the estimator must inherit from the `CallbackSupportMixin` class.\nclass SimpleKMeans(CallbackSupportMixin, BaseEstimator): # noqa: F811\n def __init__(self, n_clusters=6, n_iter=100, random_state=None):\n self.n_clusters = n_clusters\n self.n_iter = n_iter\n self.random_state = random_state\n\n def _compute_labels(self, X):\n return np.argmin(euclidean_distances(X, self.cluster_centers_), axis=1)\n\n # Then the `fit` function must be decorated with the `with_callbacks`\n # decorator, which takes care of the proper teardown of callbacks.\n @with_callbacks\n def fit(self, X, y=None):\n X = validate_data(self, X)\n random_state = check_random_state(self.random_state)\n # The `CallbackContext` object must be instantiated with the\n # `_init_callback_context` method provided by the mixin, which calls the\n # `setup` hooks of the callbacks. This context corresponds to the root task of\n # the fit function.\n callback_ctx = self._init_callback_context(\n task_name=\"fit\", max_subtasks=self.n_iter\n )\n # Then the callback context's `call_on_fit_task_begin` method must be called. It\n # will call all the callbacks' `on_fit_task_begin` hooks. The `estimator`\n # argument is mandatory and optional `kwargs` can be passed to provide extra\n # contextual information for the callbacks, for example here `X` and `y` are\n # passed. See the following note for more details on these extra `kwargs`.\n callback_ctx.call_on_fit_task_begin(estimator=self, X=X, y=y)\n\n self.cluster_centers_ = random_state.rand(self.n_clusters, X.shape[1])\n\n for i in range(self.n_iter):\n # For each sub-task of fit (here each iteration of the loop), a sub-context\n # must be created with the callback context's `subcontext` method.\n subcontext = callback_ctx.subcontext(task_name=\"fit iteration\")\n # The sub-context corresponds to a new sub-task, so its\n # `call_on_fit_task_begin` method must also be called.\n subcontext.call_on_fit_task_begin(estimator=self, X=X, y=y)\n\n labels = self._compute_labels(X)\n\n for k in range(self.n_clusters):\n if (labels == k).any():\n self.cluster_centers_[k] = X[labels == k].mean(axis=0)\n\n # After each sub-task, the `call_on_fit_task_end` method of its sub-context\n # must be called, also with `estimator` as a mandatory argument and optional\n # `kwargs`. It will call all the callbacks' `on_fit_task_end` hooks. Here\n # the extra `kwargs` contain a `reconstruction_attributes` callable,\n # which returns the necessary attributes to generate an estimator instance\n # ready to predict, as if the fit process just stopped at this step.\n if subcontext.call_on_fit_task_end(\n estimator=self,\n X=X,\n y=y,\n reconstruction_attributes=lambda: {\n \"cluster_centers_\": self.cluster_centers_,\n },\n ):\n # The `call_on_fit_task_end` method returns a boolean, which is set\n # to True if any of the callbacks' `on_fit_task_end` methods return\n # True. This enables the interruption of the `fit` process by the\n # callbacks, for example to implement early stopping. Thus the\n # `call_on_fit_task_end` method can be used in an `if` / `break` block\n # to enable such interruptions.\n break\n\n # After the root task of the fit function is done, the `call_on_fit_task_end`\n # method of its callback context must be called.\n callback_ctx.call_on_fit_task_end(\n estimator=self,\n X=X,\n y=y,\n reconstruction_attributes=lambda: {\n \"cluster_centers_\": self.cluster_centers_,\n },\n )\n\n # The callbacks' `teardown` hooks are called automatically in the decorator,\n # after fit finishes, even if it crashed.\n return self\n\n def predict(self, X):\n check_is_fitted(self)\n return self._compute_labels(X)\n\n def transform(self, X):\n check_is_fitted(self)\n return euclidean_distances(X, self.cluster_centers_)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "

Note

See the documentation of the methods\n :meth:`~CallbackContext.call_on_fit_task_begin` and\n :meth:`~CallbackContext.call_on_fit_task_end` for the description of the `kwargs`\n they can accept. These `kwargs` are optional, but an estimator should provide all\n the ones it is capable of producing in each task to be compatible with a maximum\n number of callbacks.

\n\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Registering callbacks to the custom estimator\nNow the `SimpleKMeans` estimator can be used with callbacks, for example with the\n:class:`~ProgressBar` callback to monitor progress.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "estimator = SimpleKMeans(random_state=rng)\ncallback = ProgressBar()\nestimator.set_callbacks(callback)\nestimator.fit(X)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Custom meta-estimator\nNow we demonstrate how to implement a custom meta-estimator that supports callbacks.\nFor the example, we implement a simplified version of a grid search, where only a list\nof parameter combinations is searched through instead of a grid, parallelizing the\nevaluation of the parameters.\nLet's start with the implementation without the callback support.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Function to run in parallel, it fits and scores an estimator on the folds of a CV.\ndef _fit_and_score_cv(estimator, X, y, cv, score_func):\n scores_per_fold = []\n # We iterate over the folds of the CV split.\n for train_idx, test_idx in cv.split(X):\n # A clone of the estimator is used for the current fold.\n cloned_estimator = clone(estimator)\n # The split of the current fold is applied to the data.\n X_train, X_test = X[train_idx], X[test_idx]\n y_train, y_test = (y[train_idx], y[test_idx]) if y is not None else (None, None)\n # The clone of the estimator is fitted.\n cloned_estimator.fit(X_train, y_train)\n # Its score is computed.\n scores_per_fold.append(score_func(cloned_estimator, X_test, y_test))\n return scores_per_fold\n\n\nclass SimpleGridSearch(BaseEstimator):\n def __init__(self, estimator, param_list, cv, score_func, n_jobs=1):\n # the estimator to evaluate\n self.estimator = estimator\n # the list of parameter combinations to iterate over\n self.param_list = param_list\n # the number of splits for the CV, or a CV splitter instance\n self.cv = cv\n # the scoring function\n self.score_func = score_func\n # number of jobs for parallelization\n self.n_jobs = n_jobs\n\n def fit(self, X, y=None):\n # We use a cross-validator instance to evaluate each parameter combination on\n # multiple folds.\n cv = check_cv(self.cv)\n\n # We iterate over the parameter combinations in parallel, fitting an estimator\n # and computing a score value for each fold.\n scores_per_fold = Parallel(n_jobs=self.n_jobs)(\n delayed(_fit_and_score_cv)(\n estimator=clone(self.estimator).set_params(**params),\n X=X,\n y=y,\n cv=cv,\n score_func=self.score_func,\n )\n for params in self.param_list\n )\n\n # The `cv_results_` attribute holds the score values for each parameter\n # combination and fold, as a list of tuples, each one of the form\n # (parameter combination, list of scores per fold).\n self.cv_results_ = list(zip(self.param_list, scores_per_fold))\n\n return self" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now let's update the class to support callbacks.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# The parallelized function needs to receive the callback context corresponding to its\n# task and the instance calling it.\ndef _fit_and_score_cv(estimator, X, y, cv, score_func, outer_subcontext, caller):\n # The outer sub-context's `call_on_fit_task_begin` must be called.\n outer_subcontext.call_on_fit_task_begin(estimator=caller, X=X, y=y)\n scores_per_fold = []\n for i, (train_idx, test_idx) in enumerate(cv.split(X)):\n cloned_estimator = clone(estimator)\n X_train, X_test = X[train_idx], X[test_idx]\n y_train, y_test = (y[train_idx], y[test_idx]) if y is not None else (None, None)\n # For each inner iteration a sub-context must be created.\n inner_subcontext = outer_subcontext.subcontext(task_name=f\"fold {i}\")\n # Since a sub-estimator is fitted in this task, the callbacks must be propagated\n # to that estimator with the `propagate_callback_context` context manager. Note\n # that only the callbacks following the `AutoPropagatedCallback` protocol can be\n # propagated.\n with inner_subcontext.propagate_callback_context(cloned_estimator):\n # After the propagation, the inner sub-context's `call_on_fit_task_begin`\n # method must be called.\n inner_subcontext.call_on_fit_task_begin(\n estimator=caller, X=X_train, y=y_train\n )\n cloned_estimator.fit(X_train, y_train)\n scores_per_fold.append(score_func(cloned_estimator, X_test, y_test))\n # The inner sub-context's `call_on_fit_task_end` method must be called.\n inner_subcontext.call_on_fit_task_end(\n estimator=caller, X=X_train, y=y_train\n )\n # The outer sub-context's `call_on_fit_task_end` method must be called.\n outer_subcontext.call_on_fit_task_end(estimator=caller, X=X, y=y)\n return scores_per_fold\n\n\n# The class must inherit from `CallbackSupportMixin`.\nclass SimpleGridSearch(CallbackSupportMixin, BaseEstimator): # noqa: F811\n def __init__(self, estimator, param_list, cv, score_func, n_jobs=1):\n self.estimator = estimator\n self.param_list = param_list\n self.cv = cv\n self.score_func = score_func\n self.n_jobs = n_jobs\n\n # The `fit` method must be decorated.\n @with_callbacks\n def fit(self, X, y=None):\n cv = check_cv(self.cv)\n # The callback context must be instantiated, which also calls the `setup` hooks\n # of the callbacks.\n callback_ctx = self._init_callback_context(\n task_name=\"fit\", max_subtasks=len(self.param_list)\n )\n # The `call_on_fit_task_begin` method of this context must be called.\n callback_ctx.call_on_fit_task_begin(estimator=self, X=X, y=y)\n\n # The sub-tasks of the `fit` function are nested on two levels : the outer\n # iterations over parameter combinations and the inner iterations over CV folds.\n # Sub-contexts must be created for each of these levels. For the outer level,\n # the sub-contexts are instantiated outside of the parallelized function. In\n # order to prevent any racing condition during their creation, these\n # sub-contexts must be all created before the parallelization.\n outer_subcontexts = [\n callback_ctx.subcontext(\n task_name=\"param iteration\", max_subtasks=cv.get_n_splits()\n )\n for _ in range(len(self.param_list))\n ]\n\n scores_per_fold = Parallel(n_jobs=self.n_jobs)(\n delayed(_fit_and_score_cv)(\n estimator=clone(self.estimator).set_params(**params),\n X=X,\n y=y,\n cv=cv,\n score_func=self.score_func,\n outer_subcontext=outer_subcontexts[i],\n caller=self,\n )\n for i, params in enumerate(self.param_list)\n )\n\n self.cv_results_ = list(zip(self.param_list, scores_per_fold))\n\n # The root context's `call_on_fit_task_end` must be called.\n callback_ctx.call_on_fit_task_end(estimator=self, X=X, y=y)\n\n # The callbacks' `teardown` hooks are called automatically in the decorator.\n return self" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The main difference with a simple estimator is that the callbacks must be propagated\nto the sub-estimators through the corresponding callback sub-context's\n:meth:`~CallbackContext.propagate_callback_context` context manager.\n\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "

Note

A meta-estimator that supports callback can be used with sub-estimators that do\n not. In that case a warning is raised when trying to propagate the callbacks\n and the callbacks are ignored in the sub-estimator.

\n\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Registering callbacks to the meta-estimator\nCallbacks are registered to a meta-estimator the same way as to regular estimators.\nThe callbacks which respect the :class:`~AutoPropagatedCallback` protocol (such as\n:class:`~ProgressBar`) are propagated to the sub-estimators.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "param_list = [{\"n_clusters\": 5, \"n_iter\": 20}, {\"n_clusters\": 4, \"n_iter\": 50}]\n\n\ndef score_func(estimator, X, y=None):\n return np.sum(estimator.transform(X).min(axis=1))\n\n\nsub_estimator = SimpleKMeans(random_state=rng)\nmeta_estimator = SimpleGridSearch(\n estimator=sub_estimator, param_list=param_list, cv=4, score_func=score_func\n)\ncallback = ProgressBar()\nmeta_estimator.set_callbacks(callback)\nmeta_estimator.fit(X)" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.14.6" } }, "nbformat": 4, "nbformat_minor": 0 }