Capstone: The Pipeline

Step two. We connect the stages into a pipeline, so that one command runs everything, in the right order, and records the results. The training script also logs to the experiment tracker, so that every run is remembered. Let's read the training script, and then run the pipeline.

• Three stages, one command
• Training logs to the tracker
• Reproducible by design

The training script

Here is the training script. It reads the settings from the parameters file, and the data from the data folder. It splits the data, keeping the same share of leavers in both parts. Then it opens a tracking run. It trains a random forest, with balanced class weights, because leavers are rare. It computes the F one score, on the held out part, and logs the parameters, and the score, to the tracker. Finally, it saves the model to a file, and the score to a small metrics file, which the pipeline tracks. No magic, just the pieces you already know.

import json
import joblib
import mlflow
import pandas as pd
import yaml
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import f1_score
from sklearn.model_selection import train_test_split

params = yaml.safe_load(open("params.yaml"))["train"]
df = pd.read_csv("data/churn.csv")
X, y = df.drop(columns="churn"), df["churn"]
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.25, random_state=0, stratify=y)

mlflow.set_tracking_uri("sqlite:///mlflow.db")
mlflow.set_experiment("churn-capstone")
with mlflow.start_run():
    model = RandomForestClassifier(**params, class_weight="balanced", random_state=0).fit(Xtr, ytr)
    f1 = f1_score(yte, model.predict(Xte))
    mlflow.log_params(params)
    mlflow.log_metric("f1", f1)

joblib.dump(model, "models/model.joblib")
json.dump({"f1": round(f1, 3)}, open("metrics.json", "w"))
print("f1", round(f1, 3))

Run the pipeline

Now we run the pipeline with one command. D V C runs the three stages in order. Prepare writes the data. Validate checks it, and the data is valid. Train fits the model, and prints the F one score. Then we show the tracked metric, and the graph of the pipeline. The pipeline is a small diamond, prepare feeds both validation, and training. If the data were invalid, the validate stage would fail, and we would see it right here.

dvc repro 2>&1 | grep -E 'Running|valid|^f1'
Running stage 'prepare':
Running stage 'validate':
> python src/validate.py
data valid: 4000 rows
Running stage 'train':
f1 0.389
dvc metrics show
Path          f1
metrics.json  0.389
dvc dag
          +---------+
          | prepare |
          +---------+
           *        **
         **           **
        *               *
+----------+         +-------+
| validate |         | train |
+----------+         +-------+

Summary