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

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

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

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

Исходный код

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

from src import train

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

from my_classifier import train

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

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

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

• metrics.py ^[8] содержит функции для вычисления метрик (в том числе конспект ^[9].

• Кроме того, могут существовать preprocess.py ^[10] для специфической предварительной обработки (например, выделение спектрограммы для аудиосигнала; или токенизация для задач NLP) и postprocess.py ^[11] (например, некоторые типы декодеров).Пре-/постпроцесс идут, соответственно, до и после основной обучаемой модели.

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

• data/ содержит модуль dataset.py ^[12] для батч-генерации (загрузка с диска или 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. Думаю, они не нуждаются в пояснениях.

Подытожу

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