The model trainer

The model trainer is the stage that fits a model on the transformed training matrix, scores it on the held-out test rows, and writes the fitted model and its metrics into the artifacts folder as two files.

Last updated: 19 Oct, 2026 · MLOps

Data transformation with ColumnTransformer handed this stage four things: a training matrix, its targets, a test matrix and its targets. The stage fits, scores and saves. Everything that makes it a pipeline stage rather than a notebook cell is in those last two words, because the metrics file is what the experiment tracking, the DVC and the retraining lessons all read.

Reading the model's settings from params.yaml

The forest's size and depth sit in the same file as the split settings. A different depth is then a config change that shows up in a diff, not an edited line somewhere inside a class.

TARGET_COLUMN: math_score
TEST_SIZE: 0.2
RANDOM_STATE: 42
N_ESTIMATORS: 100
MAX_DEPTH: 12

The clip finishes the pipeline with LogisticRegression; this stage finishes it with RandomForestRegressor, because the running example predicts a maths score rather than a class.

Fitting the forest

A random forest regressor fits a hundred trees and averages them. Every setting it takes comes from the config, including the seed, so two runs of the same config produce the same model.

from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

class ModelTrainer:
    def __init__(self, config: ModelTrainerConfig):
        self.config = config

    def train(self, Xtr, ytr, Xte, yte) -> dict:
        # every setting comes from params.yaml, the seed included
        model = RandomForestRegressor(
            n_estimators=self.config.n_estimators,
            max_depth=self.config.max_depth,
            random_state=self.config.random_state,
        )
        model.fit(Xtr, ytr)
        pred = model.predict(Xte)

Scoring r2, MAE and RMSE

Three numbers, rounded and saved as a dict. Rounding here is deliberate: a metrics file compared between runs should not differ in the twelfth decimal place because of a floating-point detail.

        # rounded, so two runs' metrics files differ only where the model does
        metrics = {
            "r2": round(float(r2_score(yte, pred)), 4),
            "mae": round(float(mean_absolute_error(yte, pred)), 4),
            "rmse": round(float(mean_squared_error(yte, pred) ** 0.5), 4),
        }

        save_object(self.config.model_file, model)
        save_json(self.config.metrics_file, metrics)
        return metrics

Note the last line. mean_squared_error used to take squared=False and return the root for you; scikit-learn removed the argument, so the root is taken in the code with ** 0.5.

The argument that no longer exists

Copying the older call into scikit-learn 1.9.1 gives this, which is worth meeting once so the traceback is recognisable.

# the argument older tutorials pass to get the root back directly
print(mean_squared_error(yte, pred, squared=False))

Traceback (most recent call last):
  File "main.py", line 2, in <module>
    print(mean_squared_error(yte, pred, squared=False))
TypeError: got an unexpected keyword argument 'squared'

The fix is the line the component already uses: take the square root yourself.

Training the model and reading the metrics

The same fit, flat, on the matrices the transformation stage produced. The seed is set, so these three numbers come back on any machine with these versions.

from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

model = RandomForestRegressor(n_estimators=100, max_depth=12, random_state=42)
model.fit(Xtr_t, ytr)
pred = model.predict(Xte_t)

# RMSE is taken by hand: mean_squared_error no longer returns the root for you
metrics = {"r2": round(float(r2_score(yte, pred)), 4),
           "mae": round(float(mean_absolute_error(yte, pred)), 4),
           "rmse": round(float(mean_squared_error(yte, pred) ** 0.5), 4)}
print(metrics)

{'r2': 0.8485, 'mae': 4.6996, 'rmse': 6.0721}

• r2 is 0.8485, which says the model accounts for about 85 percent of the variation in the test rows' maths scores. It has no units, so it is the number to compare across datasets.
• MAE is 4.6996, in marks: the average prediction is off by a little under five marks. This is the number to quote to somebody who has to act on the prediction.
• RMSE is 6.0721, also in marks, and larger than the MAE because squaring the errors before averaging weighs the big misses more heavily. The gap between the two is a hint about how often the model is badly wrong rather than slightly wrong.

Saving the model and its metrics file

Two writes end the stage. The model goes to disk as a joblib file, and the metrics go to disk as JSON, which anything can read without importing scikit-learn.

import json
import os
import joblib

os.makedirs("artifacts/model_trainer", exist_ok=True)
joblib.dump(model, "artifacts/model_trainer/model.joblib")

# the metrics file is the one a later run compares against, so it is written as data
json.dump(metrics, open("artifacts/model_trainer/metrics.json", "w"), indent=4)
print(open("artifacts/model_trainer/metrics.json").read())

{
    "r2": 0.8485,
    "mae": 4.6996,
    "rmse": 6.0721
}

• The JSON file is a tracked output, small enough to commit and easy to diff, so two runs can be compared by reading two files.
• The model file sits next to it under the same run's folder, which keeps a model and the score it earned together.
• Nothing in the file names the run, which is the gap experiment tracking fills.

r2 against MAE and RMSE

Where the metrics file is read later

• Comparing two runs, where a second configuration writes its own metrics and the two files are read side by side.
• A pipeline tool, where the file is declared as the stage's metrics output and shown without any code of yours.
• A retraining decision, where a live model's score drifting below the one in this file is the trigger.

Related

• Previous: Data transformation with ColumnTransformer
• Next: The prediction pipeline
• Reference: RandomForestRegressor in the scikit-learn documentation