На Github существует множество ML-проектов. Большинство из них предоставляют скрипты для обучения, тестирования, вывода моделей. Но почти все они организованы по-разному. Иногда неясно, как запустить этап конвейера, как подготовить данные или какие модели используются для предсказаний. Более того, когда разработчик заглядывает в чужой проект, он тратит много времени на то, чтобы разобраться в структуре.

В этом посте я расскажу о шаблоне ML-проекта на основе CookieCutter на примере задачи классификации. Но вообще такой шаблон может быть использован для решения множества других ML-задач. Ссылка на проект для примера.

Итак, давайте создадим наш проект шаг за шагом.

Исходный код

Как вы знаете, все ML-проекты состоят из трех основных компонентов: train, test, inference (иногда еще называют эксплуатация или предсказания). Наша первая папка с python-кодом называется src в соответствии с CookieCutter. Но здесь есть нюанс. Конечно, src — это вроде бы типичное название для папки с исходным кодом (не только для ML-проектов). И если вы собираетесь использовать эту папку в качестве python-пакета, то она будет импортирована следующим образом

from src import train

Но что такое src в нашем скрипте? Что мы собираемся здесь тренировать? Не понятно. Поэтому имеет смысл переименовать src в <имя_проекта>.

from my_classifier import train

Если вы знаете, как оставить папку с именем src, но при этом, чтоб модуль был доступен под другим именем, подскажите в комментариях:)

В нашем пакете есть три основных модуля для запуска обучения, тестирования (или валидации) и инференса: train.py, test.py, inference.py. Они должны иметь CLI для запуска из консоли или создания конвейеров (например, DVC pipelines). Я рекомендую добавить в inference.py обе возможности: предсказание для одного входного файла и для предсказаний набора файлов в директории.

К этим трём основным добавляются еще:

metrics.py содержит функции для вычисления метрик (в том числе лосс‑функции) во время тестирования (валидации) и обучения. В моей задаче классификации это FPR, TPR, F-score, ROC-кривые, PR-кривые, Cross Entropy и что-то еще. Если этот код будет раздуваться, то разумно будет выделить их отдельно в подмодуль метрик. Подробнее про метрики классификации делал такой конспект.
Кроме того, могут существовать preprocess.py для специфической предварительной обработки (например, выделение спектрограммы для аудиосигнала; или токенизация для задач NLP) и postprocess.py (например, некоторые типы декодеров).Пре-/постпроцесс идут, соответственно, до и после основной обучаемой модели.

Помимо этих скриптов есть подмодули:

data/ содержит модуль dataset.py для батч-генерации (загрузка с диска или S3, препроцессинг, перемешивание формирования батча). Также модуль make_dataset.py содержит функции для преобразования данных, такие как очистка, удаление дубликатов, удаление тишины и другие.
models/ содержит модули для каждой модели, которая используется в этом проекте (CNN, ResNet, Transformer). Да, для решения одной и той же задачи у вас могут быть разные модели. Главное, чтоб их вход и выход соответствовали единому формату.
Хорошей практикой является определение родительского класса для всех моделей, который описывает формат ввода/вывода данных, а также некоторые дополнительные методы (в частности, сохранение, загрузка, инициализация весов). Если все модели соответствуют этому формату, замена одной модели на другую не приведет к нарушению работы. Пример реализации абстрактного класса BaseModel.

import torch
from typing import Tuple
from abc import ABCMeta, abstractmethod


class BaseModel(torch.nn.Module, metaclass=ABCMeta):
   def __init__(self, *args, **kwargs):
       super().__init__()

   def init_params(self, *args, **kwargs):
       """
       Set specific weights initialization
       """
       raise NotImplementedError()
   
   @abstractmethod
   def forward(self, x:torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
       """
       B - batch size
       T - time dim
       F - feature dim
       C - n classes
       Args:
           x (B, F, T): input features
       Returns:
           tuple
             (B, C): output logits
             (B, C): output probs (output of softmax)
       """
       pass
         
   def load(self, weights) -> None:
       state_dict = torch.load(weights)
       self.load_state_dict(state_dict)

   def save(self, weights_path: str):
       torch.save(self.state_dict(), weights_path)

utils/ хранит все дополнительные инструменты, например, логгер, менеджер (или треккер) экспериментов, hypertune и другие.

Другие составляющие проекта

С исходными файлами разобрались, но ML-проект состоит не только из них.

data/ - папка с наборами данных: файлами с метками и файлами содержимого (изображения, аудиозаписи, тексты). Согласно CookieCutter data/ имеет четыре вложенные папки: raw/ для исходных данных, interim/ для промежуточных этапов обработки, processed/ для данных, готовых для обучения и тестирования, external/ данные с внешних источников (3rd party).
experiments/ - хранит результаты обучения и тестирования (метрики, логи, конфиги, веса). В общем все, что нужно для оценки эксперимента, сравнений друг с другом, а также воспроизведения.

docker/ — все, что нужно для создания и запуска контейнеров docker. Например Dockerfile.
docs/ — некоторые конкретные документы и спецификации (pdf или markdown), описывающие проект. Как пользоваться CLI, ноутбуками, как импортировать проект, как готовить данные, как управлять конфигом.
models/ (или weights/) — папка с моделями, в которой я предпочитаю хранить веса, пригодные для production. А у вас может быть какое-нибудь свое registry. Желательно, чтоб модели, хранимые здесь, имели свое описание (метрики, скорость, GPU memory usage) в отдельном документе.
notebook/ — в некоторых случаях вам может потребоваться тетрадки jupyter. Например, для быстрого анализа данных (построения графиков и т.п.)
references/ — если ваш проект ссылается на какие-то конкретные библиотеки или внешние репозитории, их следует загрузить сюда. Например, ваш NLP-проект использует токенизатор из HuggingFace. Или реализации моделей из другого репозитория, которые вы не копируете к себе, а импортируете. Чтоб не раздувать репозиторий, лучше эту папку кинуть в .gitignore.
reports/ — сохраняйте здесь отчеты о своих исследованиях. Например, месяц вашей работы привел к появлению новых моделей, которые вы сравнили с предыдущими на графиках и т.п. и сформировали pdf отчет. Это поможет вам (или менеджерам проекта) оценить результат.
tests/ — скрипты в этой папке предназначены для тестирования исходного кода. Можно настроить работу так, чтобы они запускались автоматически перед слиянием с основной веткой. Скажем, если вы обновили модель, вам стоит удостовериться, что формат не поменялся, что inference работает, как ожидается. Это уменьшит шанс появления багов.
config.yml — один из многих способов хранить конфигурацию проекта. Я привык держать все настройки проекта в одном файле. Хотя можно хранить такие файлы отдельно для моделей, обучения, данных и т. д. В этом случае предложил бы завести папку configs, в которой будут эти файлы. Стоит сказать, что конфигурация — очень индивидуальный момент. Например, learning rate можно хранить в конфиг. файле, а можно передавать в качестве доп. аргумента при запуске обучения. Кому как.
Также необходимо добавить README.md и requirements.txt. Думаю, они не нуждаются в пояснениях.

Подытожу

Все описанное является лишь шаблоном, на базе которого может быть проще построить ваш проект. Я не призываю использовать каждый из описанных компонентов, поскольку некоторые из них могут быть для вас избыточными. А если у вас есть свои наработанные практики, можете поделиться ими в комментариях :-)

Организация ML-проекта с примерами