Experiment Tracking

Experiment Tracking#

Important

Experiment tracking requires ploomber-engine 0.0.11 or higher

ploomber-engine comes with a tracker that allows you to log variables, and plots without modifying your code. A common use case is for managing Machine Learning experiments. However, you can use it for any other computational experiments where you want to record input parameters, outputs, and analyze results.

Under the hood, it uses the SQliteTracker from our sklearn-evaluation package, this means that you will be able to query, and aggregate your experiments with SQL!

For this example, we’ll train several Machine Learning models and evaluate their performance.

Dependencies#

Before we run this example, we’ll need to install the required dependencies:

pip install ploomber-engine sklearn-evaluation jupytext --upgrade

Now, let’s download the sample script we’ll use, Jupyter notebooks (.ipynb) are also supported:

%%sh
curl -O https://raw.githubusercontent.com/ploomber/posts/master/experiment-tracking/fit.py

  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current

                                 Dload  Upload   Total   Spent    Left  Speed

  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0
100  1718  100  1718    0     0   8354      0 --:--:-- --:--:-- --:--:--  8380

Let’s create our grid of parameters, we’ll execute one experiment for each combination:

from ploomber_engine.tracking import track_execution
from sklearn.model_selection import ParameterGrid

grid = ParameterGrid(
    dict(
        n_estimators=[10, 15, 25, 50],
        model_type=["sklearn.ensemble.RandomForestClassifier"],
        max_depth=[5, 10, None],
        criterion=["gini", "entropy"],
    )
)

Let’s now execute the script multiple times, one per set of parameters, and store the results in the experiments.db SQLite database:

for idx, p in enumerate(grid):
    if (idx + 1) % 10 == 0:
        print(f"Executed {idx + 1}/{len(grid)} so far...")

    track_execution("fit.py", parameters=p, quiet=True, database="experiments.db")

Executed 10/24 so far...

Executed 20/24 so far...

If you prefer so, you can also execute scripts from the terminal:

python -m ploomber_engine.tracking fit.py \
    -p n_estimators=100,model_type=sklearn.ensemble.RandomForestClassifier,max_depth=10

No extra code needed#

If you look at the script’s source code, you’ll see that there are no extra imports or calls to log the experiments, it’s a vanilla Python script.

The tracker runs the code, and it logs everything that meets any of the following conditions:

1. A line that contains a variable name by itself

accuracy = accuracy_score(y_test, y_pred)
accuracy

2. A line that calls a function with no variable assignment

plot.confusion_matrix(y_test, y_pred)

Note that logs happen in groups (not line-by-line). A group is defined as a contiguous set of lines delimited with line breaks:

# first group
accuracy = accuracy_score(y_test, y_pred)
accuracy

# second group
plot.confusion_matrix(y_test, y_pred)

Regular Python objects are supported (numbers, strings, lists, dictionaries), and any object with a rich representation will also work (e.g., matplotlib charts).

Querying experiments with SQL#

After finishing executing the experiments, we can initialize our database (experiments.db) and explore the results:

from sklearn_evaluation import SQLiteTracker

tracker = SQLiteTracker("experiments.db")

Since this is a SQLite database, we can use SQL to query our data using the .query method. By default, we get a pandas.DataFrame:

df = tracker.query(
    """
SELECT *
FROM experiments
LIMIT 5
"""
)

type(df)

pandas.core.frame.DataFrame

df.head(3)

	created	parameters	comment
uuid
3f34c993	2023-09-13 02:08:53	{"criterion": "gini", "max_depth": 5, "model_t...	None
39632d5b	2023-09-13 02:08:54	{"criterion": "gini", "max_depth": 5, "model_t...	None
7ed51791	2023-09-13 02:08:55	{"criterion": "gini", "max_depth": 5, "model_t...	None

The logged parameters are stored in the parameters column. This column is a JSON object, and we can inspect the JSON keys with the .get_parameters_keys() method:

tracker.get_parameters_keys()

['accuracy',
 'classification_report',
 'confusion_matrix',
 'criterion',
 'f1',
 'max_depth',
 'model_type',
 'n_estimators',
 'precision',
 'precision_recall',
 'recall']

To get a sample query, we can use the .get_sample_query() method:

print(tracker.get_sample_query())

SELECT
    uuid,
    json_extract(parameters, '$.accuracy') as accuracy,
    json_extract(parameters, '$.classification_report') as classification_report,
    json_extract(parameters, '$.confusion_matrix') as confusion_matrix,
    json_extract(parameters, '$.criterion') as criterion,
    json_extract(parameters, '$.f1') as f1,
    json_extract(parameters, '$.max_depth') as max_depth,
    json_extract(parameters, '$.model_type') as model_type,
    json_extract(parameters, '$.n_estimators') as n_estimators,
    json_extract(parameters, '$.precision') as precision,
    json_extract(parameters, '$.precision_recall') as precision_recall,
    json_extract(parameters, '$.recall') as recall
    FROM experiments
LIMIT 10

Let’s perform a simple query that extracts our metrics, sorts by F1 score, and returns the top three experiments:

tracker.query(
    """
SELECT
    uuid,
    json_extract(parameters, '$.f1') as f1,
    json_extract(parameters, '$.accuracy') as accuracy,
    json_extract(parameters, '$.criterion') as criterion,
    json_extract(parameters, '$.n_estimators') as n_estimators,
    json_extract(parameters, '$.precision') as precision,
    json_extract(parameters, '$.recall') as recall
FROM experiments
ORDER BY f1 DESC
LIMIT 3
""",
    as_frame=False,
    render_plots=False,
)

uuid	f1	accuracy	criterion	n_estimators	precision	recall
89547f24	0.806812	0.800606	entropy	50	0.778912	0.836784
55f228de	0.803288	0.796970	gini	50	0.775510	0.833130
0c5a0d1f	0.799298	0.792121	entropy	25	0.769144	0.831912

Rendering plots#

Plots are essential for evaluating our analysis, and the experiment tracker can log them automatically. So let’s pull the top 2 experiments (by F1 score) and get their confusion matrices, precision-recall curves, and classification reports. Note that we passed as_frame=False and render_plots=True to the .query method:

top_two = tracker.query(
    """
SELECT
    uuid,
    json_extract(parameters, '$.f1') as f1,
    json_extract(parameters, '$.confusion_matrix') as confusion_matrix,
    json_extract(parameters, '$.precision_recall') as precision_recall,
    json_extract(parameters, '$.classification_report') as classification_report
FROM experiments
ORDER BY f1 DESC
LIMIT 2
""",
    as_frame=False,
    render_plots=True,
)

top_two

uuid	f1	confusion_matrix	precision_recall	classification_report
89547f24	0.806812
55f228de	0.803288

Comparing plots#

If you want to zoom into a couple of experiment plots, you can select one and create a tab view. First, let’s retrieve the confusion matrices from the top 2 experiments:

top_two.get("confusion_matrix")

The buttons in the top bar allow us to switch between the two plots of the selected experiments. We can do the same with the precision-recall curve:

top_two.get("precision_recall")

Note that the tabs will take the value of the first column in our query (in our case, the experiment ID). So let’s switch them to show the F1 score via the index_by argument:

top_two.get("classification_report", index_by="f1")

Aggregation and plotting#

Using SQL gives us power and flexibility to answer sophisticated questions about our experiments.

For example, let’s say we want to know the effect of the number of trees (n_estimators) on our metrics. We can quickly write a SQL that aggregates by n_estimators and computes the mean of all our metrics; then, we can plot the values.

First, let’s run the query. Note that we’re passing as_frame=True so we get a pandas.DataFrame we can manipulate:

df = tracker.query(
    """
SELECT
    json_extract(parameters, '$.n_estimators') as n_estimators,
    AVG(json_extract(parameters, '$.accuracy')) as accuracy,
    AVG(json_extract(parameters, '$.precision')) as precision,
    AVG(json_extract(parameters, '$.recall')) as recall,
    AVG(json_extract(parameters, '$.f1')) as f1
FROM experiments
GROUP BY n_estimators
""",
    as_frame=True,
).set_index("n_estimators")

df

	accuracy	precision	recall	f1
n_estimators
10	0.748636	0.725184	0.807349	0.761826
15	0.762374	0.730451	0.829984	0.776724
25	0.765657	0.731813	0.839017	0.781270
50	0.773081	0.738236	0.846833	0.788292

Now, let’s create the plot:

import matplotlib.pyplot as plt

fig, ax = plt.subplots()

for metric in ["accuracy", "f1", "recall", "precision"]:
    df[metric].plot(ax=ax, marker="o", style="--")

ax.legend()
ax.grid()
ax.set_title("Effect of increasing n_estimators")
ax.set_ylabel("Metric")
_ = ax.set_xticks(df.index)

../_images/058b54a993acb60af7d3a1ed81018511c846e2acf99f8983ec2b5891085efacf.png

We can see that increasing n_estimators drives our metrics up; however, after 25 estimators, the improvements get increasingly small. Considering that a larger number of n_estimators will increase the runtime of each experiment, we can use this plot to inform our following experiments, cap n_estimators, and focus on other parameters.

If you need more details on how to work with the SQLiteTracker, you cna jump directly to the full example in the documentation.