Мы в ПГК занимаемся грузоперевозками, причем решаем различные транспортные задачи не только методами математической оптимизации, но и с помощью моделей машинного обучения. Наши дата-сайентисты проводят десятки экспериментов — в том числе и без необходимости прибегать к инструментам логирования вроде MLflow. В этом им помогает компактный Python-класс. Расскажем, как он устроен, и поделимся кодом.

Мы активно используем машинное обучение для разработки решений, упрощающих жизнь наших сотрудников. Например, у нас есть «Оптимизатор ремонтов», который помогает составлять план ремонта вагонов и в то же время ранжировать депо по стоимости работ. Также мы используем «Прогнозатор», который предсказывает объёмы грузоперевозок между ж/д станциями для последующего планирования продаж.

Поскольку мы работаем с десятками различных моделей, нам необходимо вести учет переменных, управлять метриками, сравнивать прогнозы. Да, для версионирования экспериментов мы используем MLflow, который предлагает мощные специализированные возможности. Мы уже рассказывали об этом инструменте у себя в блоге. Однако его внедрение и поддержка подразумевает работу с инфраструктурой для логирования и мониторинга ML-экспериментов, а также управления кодом.

На простых задачах и небольших проектах не всегда нужно сложное решение для управления версиями. Если есть необходимые данные, а задача понятна, то прототип вполне можно «набросать» за один стандартный двухнедельный спринт. Иными словами, в контексте быстрого прототипирования, когда нужно просто проверить гипотезу, MLflow — это определённо overkill.

Всего один класс

Чтобы помочь дата-сайентистам разрабатывать прототипы и тестировать гипотезы [причем делать это локально], мы разработали компактный инструмент MLExperimentManager. Это — Python-класс, который предоставляет интерфейс для загрузки данных, обучения моделей, оценки их качества и логирования результатов экспериментов. Скрипт делает процесс исследования и сравнения различных ML-моделей более организованным и систематизированным.

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_curve, auc, confusion_matrix
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
import seaborn as sns
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

class MLExperimentManager:
    def __init__(self):
        self.data = None
        self.X_train, self.X_test, self.y_train, self.y_test = (None,) * 4
        self.model = None
        self.experiments_log = None
        self.log_path = 'experiments_log.csv'

    def load_data(self, dataset):
        self.data = dataset
        X = pd.DataFrame(self.data.data, columns=self.data.feature_names)
        y = pd.Series(self.data.target)
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        print("Data loaded successfully.")

    def train_model(self, classifier='RandomForest'):
        if classifier == 'RandomForest':
            self.model = RandomForestClassifier(n_estimators=100, random_state=42)
        elif classifier == 'LogisticRegression':
            self.model = LogisticRegression(max_iter=200, random_state=42)
        elif classifier == 'SVC':
            self.model = SVC(probability=True, random_state=42)
        else:
            raise ValueError(f"{classifier} is not supported")
        self.model.fit(self.X_train, self.y_train)
        print(f"Model training completed using {classifier}.")

    def evaluate_model(self, metric='accuracy'):
        predictions = self.model.predict(self.X_test)
        if metric == 'accuracy':
            result = accuracy_score(self.y_test, predictions)
        elif metric == 'precision':
            result = precision_score(self.y_test, predictions, average='macro')
        elif metric == 'recall':
            result = recall_score(self.y_test, predictions, average='macro')
        elif metric == 'f1':
            result = f1_score(self.y_test, predictions, average='macro')
        else:
            raise ValueError(f"{metric} is not a supported metric")
        print(f"Model {metric}: {result}")
        return result

    def log_experiment_to_csv(self, experiment_id, model_params, data_path, result, metric):
        if self.experiments_log is None:
            self.experiments_log = pd.DataFrame({
            'ExperimentID': experiment_id,
            'Comment': 'just_test',
            'Model': self.model.__class__.__name__,
            'ModelParams': str(model_params),
            'DataPath': data_path,
            'Metric': metric,
            'Result': result
        }, index=[0])
        else: 
            self.experiments_log = pd.concat([self.experiments_log, 
                                             pd.DataFrame({
                                            'ExperimentID': experiment_id,
                                            'Comment': 'just_test',
                                            'Model': self.model.__class__.__name__,
                                            'ModelParams': str(model_params),
                                            'DataPath': data_path,
                                            'Metric': metric,
                                            'Result': result
                                            }, index=[self.experiments_log.index[-1]+1])], ignore_index=True
                                            )
            
        log_df = pd.DataFrame(self.experiments_log)
        log_df.to_csv(self.log_path, index=False)
        print(f"Experiment {experiment_id} logged successfully to {self.log_path}.")
   
    # актуально только для бинарной классификации
    def plot_roc_curves(self):
        if hasattr(self.model, "predict_proba"):
            probas_ = self.model.predict_proba(self.X_test)
            fpr, tpr, thresholds = roc_curve(self.y_test, probas_[:, 1])
            roc_auc = auc(fpr, tpr)

            plt.figure()
            plt.plot(fpr, tpr, color='darkorange', lw=2, label='ROC curve (area = %0.2f)' % roc_auc)
            plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
            plt.xlim([0.0, 1.0])
            plt.ylim([0.0, 1.05])
            plt.xlabel('False Positive Rate')
            plt.ylabel('True Positive Rate')
            plt.title('Receiver operating characteristic example')
            plt.legend(loc="lower right")
            plt.show()
        else:
            print("The current model does not support probability predictions.")

    def plot_results(self):
        # Предсказываем классы на тестовом наборе
        y_pred = self.model.predict(self.X_test)
        
        # Считаем матрицу ошибок
        conf_mat = confusion_matrix(self.y_test, y_pred)
        
        # Визуализируем матрицу ошибок
        sns.heatmap(conf_mat, annot=True, fmt='d', cmap='Blues',
                    xticklabels=self.data.target_names, yticklabels=self.data.target_names)
        plt.xlabel('Predicted labels')
        plt.ylabel('True labels')
        plt.title('Confusion Matrix')
        plt.show()

MLExperimentManager включает несколько методов:

load_data. Метод отвечает за загрузку данных в формате Parquet/CSV и предоставляет удобный интерфейс для их предварительной обработки — например, удаления пропусков. В то же время он делит данные на обучающий и тестовый наборы. Модель позволяет проводить тесты на данных, находящихся за горизонтом обучения — то есть на выборке out-of-time (ООТ). Иными словами, её работу можно проверять на данных, которые она вообще не видела.

train_model. Метод предназначен для обучения модели. Он поддерживает самые разные алгоритмы — градиентный бустинг, случайный лес, логистическую регрессию, метод опорных векторов. Параметры выбранной модели можно оптимизировать — например, изменить максимальную глубину дерева или модифицировать количество деревьев в случайном лесе, а также поменять коэффициент регуляризации в логистической регрессии.

evaluate_model. Метод для оценки качества модели. Он поддерживает четыре метрики: accuracy, precision, recall и f1. Это — наиболее распространённые метрики для решения задач классификации, но при желании их количество можно расширить.

log_experiment_to_csv. Метод отвечает за сохранение результатов эксперимента. Они записываются в таблицу, которая хранится локально. Все записи содержат идентификатор эксперимента, комментарий, информацию об использованной модели, данные, ссылку на сэмпл, метрику, величину этой метрики на выборке out-of-time и параметры.

plot_roc_curves. Метод для визуализации фактических и предсказанных результатов. Он позволяет наглядно оценить эффективность модели, в том числе коллегам из бизнес-подразделений.

Вы можете свободно применять класс MLExperimentManager в своих проектах, выбирая необходимые параметры в методах. Лишь небольшое замечание: в текущем виде скрипт хорошо подходит для решения задач классификации. Если нужно решать задачи регрессии, код потребует доработок. Придется модифицировать функцию train_model, добавив в нее алгоритмы для решения задач регрессии, и отказаться от функции plot_roc_curves для визуализации результатов. Также потребуется изменить метрики в evaluate_model и переписать функцию plot_results под расчеты ошибок для задач регрессии (например МАЕ, МАРЕ, MSE).

Практика

Покажем, как использовать MLExperimentManager на примере знаменитого набора данных Iris из библиотеки scikit-learn. Он содержит информацию о трёх видах ирисов (Iris setosa, Iris virginica и Iris versicolor), и наша цель — обучить модель машинного обучения для классификации видов на основе четырёх характеристик: длины и ширины лепестков и чашелистиков. Начнем с подготовки окружения: импортируем необходимые библиотеки и наш класс MLExperimentManager, а также загрузим данные.

from sklearn.datasets import load_iris
from MLExperimentManager import MLExperimentManager
iris = load_iris()

Проведем инициализацию и сформируем экземпляр класса MLExperimentManager. Загрузим в него данные Iris, и скрипт автоматически разделит их на обучающую и тестовые выборки.

manager = MLExperimentManager()
manager.load_data(iris)

Далее, перейдем к тренировке модели: manager.train_model. Для примера обучим три различные модели: случайный лес (RandomForest), логистическую регрессию (LogisticRegression) и метод опорных векторов (SVC), чтобы сравнить их эффективность.

Обучаем случайный лес:

manager.train_model(classifier='RandomForest')
manager.evaluate_model(metric='accuracy')

Обучаем логистическую регрессию:

manager.train_model(classifier='LogisticRegression')
manager.evaluate_model(metric='accuracy')

Обучаем метод опорных векторов:

manager.train_model(classifier='SVC')
manager.evaluate_model(metric='accuracy')

Для оценки качества воспользуемся manager.evaluate_model с метриками accuracy, precision, recall и f1. Сразу запишем результаты в файл CSV:

accuracy = manager.evaluate_model(metric='accuracy')
manager.log_experiment_to_csv(experiment_id=1, model_params=manager.model.get_params(), data_path="iris_dataset", result=accuracy, metric='accuracy')
precision = manager.evaluate_model(metric='precision')
manager.log_experiment_to_csv(experiment_id=2, model_params=manager.model.get_params(), data_path="iris_dataset", result=precision, metric='precision')
f1 = manager.evaluate_model(metric='f1')
manager.log_experiment_to_csv(experiment_id=3, model_params=manager.model.get_params(), data_path="iris_dataset", result=f1, metric='f1')

Посмотрим, какая информация записалась в таблицу:

Используем метод plot_roc_curves для визуализации матрицы ошибок, чтобы лучше понять, как модель классифицирует различные виды ирисов:

manager.plot_results()

На графике по оси абсцисс изображены предсказанные классы моделью, а по оси ординат — реальные классы. По матрице ошибок можно понять, что на конкретном примере модель корректно провела классификацию.

Таким образом, наш MLExperimentManager позволяет легко обучить несколько моделей на наборе данных Iris и оценить их качество. И такой подход можно без проблем адаптировать для работы с другими наборами данных и моделями, что делает его ценным инструментом для исследователей и разработчиков в области машинного обучения.