Когда достаточно простого класса Python — взять и начать управлять ML-экспериментами

Мы в ПГК занимаемся грузоперевозками, причем решаем различные транспортные задачи не только методами математической оптимизации, но и с помощью моделей машинного обучения. Наши дата-сайентисты проводят десятки экспериментов — в том числе и без необходимости прибегать к инструментам логирования вроде 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 и оценить их качество. И такой подход можно без проблем адаптировать для работы с другими наборами данных и моделями, что делает его ценным инструментом для исследователей и разработчиков в области машинного обучения.