Введение

Традиционный пайплайн обучения нейронной сети

На рисунке выше изображена привычная схема пайплайна обучения нейронной сети. Условно ее можно разделить на две части:

Загрузка и подготовка данных. В экосистеме PyTorch эту роль обычно выполняет DataLoader в связке с Dataset. Вместе они решают несколько задач:

Интерфейс доступа к данным через пользовательскую реализацию Dataset
Преобразование и аугментация данных, обычно выполняемые при получении элемента в Dataset
Параллельная подготовка данных в отдельных worker-процессах
Объединение отдельных элементов в батчи с помощью collate_fn
Итеративная выдача подготовленных батчей циклу обучения

Цикл обучения нейронной сети. Реализуется вручную или с помощью стороннего фреймворка. Как правило, он включает следующие шаги:

Получение очередного батча данных от DataLoader
Прямой проход модели (forward) и получение предсказания
Вычисление функции потерь
Обратное распространение ошибки и вычисление градиентов (backward)
Обновление параметров модели оптимизатором
Дополнительно: вычисление метрик, логирование, визуализация и сохранение чекпоинтов

В такой системе можно выделить следующие проблемы:

Неявная структура данных: подготовка и преобразование данных распределены между Dataset, аугментациями, collate_fn и циклом обучения. При этом структура передаваемых данных не является явной частью пайплайна и определяется соглашениями в пользовательском коде. На практике часто используются словари со строковыми ключами, содержащие Tensor, NumPy-массивы и другие объекты.
Фрагментированный пайплайн вычислений: последовательные операции над данными выполняются разными подсистемами: часть выполняется в Dataset и DataLoader на этапе подготовки данных, другая - в цикле обучения: прямой проход модели, вычисление функции потерь, обратное распространение и обновление параметров модели.
Исполнение привязано к конкретным компонентам: DataLoader предоставляет свой механизм многопроцессорной подготовки данных, управляемый через специальный API. При этом этот механизм нельзя аналогичным образом применить к произвольной части вычислительного пайплайна.

У всех этих проблем есть одна корневая причина - пайплайн обучения не представлен как единая сущность. Его части распределены между Dataset, DataLoader, моделью и циклом обучения.

А можно ли представить весь пайплайн обучения в виде единого графа вычислений? Оказывается, можно. И сейчас я разберу один из вариантов такого подхода на базе моего фреймворка ICO, собрав полный пайплайн классификатора изображений на CIFAR-10.

Основы подхода

Пайплайн в виде единого вычислительного графа на базе ICO

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

В основе подхода лежат несколько простых принципов:

Любое вычисление представлено оператором.
Каждый оператор явно определяет тип входных и выходных данных, которые задают его сигнатуру: $I \rightarrow O$ , где - тип входных данных, - тип выходных данных.
Граф представляет собой линейный однонаправленный поток без ветвлений. Операторы соединяются в последовательную цепочку: выход одного оператора становится входом следующего: $I \rightarrow O_0 \rightarrow O_1 ...$ .
Управляющие конструкции тоже представлены операторами. Например, цикл является оператором, который многократно выполняет вложенный в него оператор. Таким образом реализуется принцип композиции - оператор может включать в себя другие операторы и выполнять их в соответствии с заложенной в него логикой.
Для сохранения состояния между операторами используется Context. Операторы могут читать и изменять его: $С_0 \rightarrow C_1 \rightarrow C_2$ .
Оператор может одновременно преобразовывать входные данные и изменять контекст: $(I_0, C_0) \rightarrow (I_1, C_1) \rightarrow ...$ . Это позволяет выразить в том же графе вычисления с состоянием, необходимые, например, для обучения модели.

Эти принципы лежат в основе архитектуры ICO, и сейчас мы приступим к реализации классификатора на CIFAR-10.

Пайплайн подготовки данных

Базовые структуры данных

Сначала определим базовые структуры данных:

CifarItem: дата-класс для одного элемента датасета CIFAR-10 с полями для изображения (image) и аннотации класса (label)
CifarDataset: класс, реализующий интерфейс доступа к CifarItem. Датасет относительно небольшой, поэтому мы загрузим его весь в память.

from collections.abc import Iterator
from dataclasses import dataclass

import torch
from torch import Tensor
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor


@dataclass(slots=True, frozen=True)
class CifarItem:
    image: Tensor  # (3, 32, 32)
    label: Tensor  # (,)


class CifarDataset:
    images: Tensor  # (N, 3, 32, 32)
    labels: Tensor  # (N,)

    def __init__(self, root: str = "data/"):
        dataset = CIFAR10(root=root, download=True, transform=ToTensor())

        # Загружаем все данные в память в тензоры
        self.images = torch.stack(
            [dataset[i][0] for i in range(len(dataset))]
        )  # (N, 3, 32, 32)
        self.labels = torch.tensor([dataset[i][1] for i in range(len(dataset))])  # (N,)

    def __len__(self) -> int:
        return len(self.images)

    def __getitem__(self, index: int) -> CifarItem:
        return CifarItem(image=self.images[index], label=self.labels[index])

    def __iter__(self) -> Iterator[CifarItem]:
        for i in range(len(self)):
            yield self[i]

Определение источника данных

В ICO все - операторы, и источник данных не исключение. Это особый вид оператора, имеющий класс IcoSource и сигнатуру $() \rightarrow Iterator[T]$ , т. e. он не принимает ничего на вход (точнее принимает None) и возвращает итератор по данным.

Тут можно сделать небольшое математическое отвлечение. С точки зрения логики у него должен быть симметричный антагонист - оператор, который принимает набор данных и ничего не возвращает . И действительно, в ICO это финальный оператор, такой как IcoSink. Он терминирует поток данных и имеет сигнатуру $Iterator[T] \rightarrow ()$ .

Тогда граф, начинающийся с источника данных и заканчивающийся финальным оператором можно рассмотреть как оператор с сигнатурой $() \rightarrow ()$ . Т.е. это самодостаточный оператор, который выполняет вычисления. Такую конструкцию можно называть замкнутым графом и именно к такому виду у нас придет пайплайн обучения нейронной сети в конце статьи, что мне кажется довольно красивым!

Итак, вернемся к определению источника данных. Для эффективности мы начнем оперировать индексами данных, а не самими данными и сделаем разделение на тренировочный и валидационный набор.

import random

from ico import IcoSource

# Создаем датасет 

dataset = CifarDataset()

# Формируем набор индексов для тренировочного и валидационного набора

dataset_size = len(dataset)
all_indices = list(range(dataset_size))
random.shuffle(all_indices)

# Используем 20% данных как валидационный набор

val_split_ratio = 0.2
train_split_ratio = 1 - val_split_ratio
train_split_size = int(dataset_size * train_split_ratio)

# Создаем источник данных для тренировки

def get_train_indices() -> Iterable[int]:
    return all_indices[: train_split_size]

train_source = IcoSource(get_train_indices, name="CIFAR10 train indices")

Вычисления в операторе ICO задаются пользовательской функцией. В случае с источником данных, это функция get_train_indices, возвращающая Iterable[T].

В ICO каждый пайплайн - это цепочка операторов, представленная в виде дерева и мы всегда можем посмотреть его план выполнения, вызвав утилиту describe():

train_source.describe()

Для отображения плана выполнения реализован рендерер, работающий с библиотекой Rich - для большей выразительности, совместимости с терминалом и просто эстетики.

Пока наш пайплайн состоит из одного оператора, поэтому в плане выполнения мы видим одну строчку. В столбце “Flow” мы видим класс оператора IcoSource, его пользовательскую функцию get_train_indices и свойство size с количеством данных. В столбце “Signature” мы видим его сигнатуру $() \rightarrow Iterator[int]$ . В последнем столбце “Name” - название, которое может задать пользователь.

Группировка в батчи

Теперь добавим случайное перемешивание индексов и группировку в батчи в наш пайплайн.

from ico import IcoBatcher


@operator()
def shuffle_indices(indices: Iterator[int]) -> Iterator[int]:
    indices_list = list(indices)
    random.shuffle(indices_list)
    yield from indices_list

batch_size = 8
batcher = IcoBatcher[int](batch_size=batch_size)

train_data_pipeline = train_source | shuffle_indices | batcher
train_data_pipeline.describe()

Для создания оператора shuffle_indices мы используем одноименную функцию с декоратором @operator(), который возвращает тип IcoOperator.

IcoBatcher это специальный оператор, который объединяет элементы в группы указанного размера и имеет сигнатуру $Iterator[T] \rightarrow Iterator[Iterator[T]]$ . В сигнатуре мы как раз видим, что появился еще один уровень вложенности , соответствующий батчу индексов.

В ICO мы соединяем операторы в цепочки с помощью оператора “или” |. Поскольку оператор типизирован с использованием дженерик параметров, статические анализаторы такие как mypy могут проверять правильное соответствие типов в таких цепочках еще до выполнения кода.

Загрузка данных

Для простоты изложения материала, сначала мы рассмотрим сценарий обработки данных в одном общем процессе, а потом реализуем многопроцессорный вариант c параллельной загрузкой.

В первом варианте как таковая загрузка тривиальна - нам надо получить элемент по индексу из датасета.


@operator()
def fetch_item(idx: int) -> CifarItem:
    return dataset[idx]

train_data_pipeline = (
    train_source 
    | shuffle_indices 
    | batcher 
    | fetch_item.stream().stream()
)
train_data_pipeline.describe()

Но погодите! Как тривиальноe чтение превратилось в два вложенных цикла и для чего конструкция .stream().stream()? Давайте разбираться!

Оператор fetch_item определен только для одного элемента, т.е. он умеет извлекать элемент датасета CifarItem по индексу. А как быть с двойной вложенностью Iterable[Iterable[int]]?

Для этого как раз есть оператор, который является управляющей конструкцией - IcoStream. Он имеет сигнатуру $Iterator[I] \rightarrow Iterator[O]$ и образует композицию с вложенным оператором body с сигнатурой $I \rightarrow O$ , который является аргументом инициализации. В своей вычислительной функции стрим итерируется по входной последовательности, применяет вложенный оператор и возвращает итератор с результатом его работы:

def _stream_fn(self, items: Iterator[I]) -> Iterator[O]:
    for item in items:
        yield self.body(item)

В плане выполнения у него есть свой рендерер, который показывает логику итерации и оператор, который применяется к каждому элементу последовательности. А .stream() это синтаксический сахар для удобства.

Значит, .stream().stream() поднимает оператор fetch_item на два уровня вложенности: внутренний IcoStream применяет его к индексам внутри батча, а внешний - к последовательности батчей. В результате сигнатура оператора становится совместимой с выходом batcher.

Формирование батча и collate_fn

Следующая задача - свернуть набор Iterator[CifarItem] в батч. Сначала опишем структуру батча данных. Она принимает набор items и складывает в тензор все изображения и все аннотации, добавляя новое измерение, соответствующее размеру батча.

@final
class CifarBatch:
    __slots__ = ("images", "labels")

    images: Tensor  # (B, 3, 32, 32)
    labels: Tensor  # (B,)

    def __init__(self, items: Iterator[CifarItem]):
        items_list = list(items)

        # Создание тензоров с изображениями и аннотациями

        self.images = torch.stack([item.image for item in items_list])  # (B, 3, 32, 32)
        self.labels = torch.stack([item.label for item in items_list])  # (B,)

Далее, определяем оператор, который будет формировать батч из набора элементов датасета, то есть делать collate.

@operator()
def collate(items: Iterator[CifarItem]) -> CifarBatch:
    batch = CifarBatch(items)
    return batch

Дополним наш пайплайн формированием батча данных:

train_data_pipeline = (
    train_source 
    | shuffle_indices 
    | batcher 
    | fetch_item.stream().stream()
    | collate.stream()
)
train_data_pipeline.describe()

Collate “схлопывает” один уровень вложенности и в последней строчке на выходе мы видим Iterator[CifarBatch] - это именно тот вход, который потребуется для дальнейшего цикла тренировки.

Аугментация данных

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

Для этого напишем небольшой набор операторов, которые применяют трансформацию к изображению у CifarItem. Реализацию трансформаций возьмем из пакета torchvision.transforms.functional.

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

from torchvision.transforms import functional as F

class CifarTransform(ABC):
    """Базовый класс для всех трансформаций с параметром вероятности применения."""

    p: float  # Вероятность применения
    factor_max_delta: float  # Максимальное значение эффекта трансформации

    def __init__(self, p: float = 1.0, factor_max_delta: float = 0.1):
        self.p = p
        self.factor_max_delta = factor_max_delta

    # Абстрактный метод для реализации трансформации

    @abstractmethod
    def _image_transform(self, image: Tensor, factor_max_delta: float) -> Tensor:
        raise NotImplementedError

    # Реализация логики применения трансформации на основе вероятности срабатывания

    def __call__(self, item: CifarItem) -> CifarItem:
        if torch.rand(1) >= self.p:
            return item

        # Получаем случайное значение фактора из равномерно распределенных значений
        # на интервале [1 - factor_max_delta: 1 + factor_max_delta) 

        factor = 1 + (torch.rand(1).item() - 0.5) * 2.0 * self.factor_max_delta
        
        return CifarItem(
            image=self._image_transform(item.image, factor), 
            label=item.label,
        )

Далее реализуем набор трансформаций для изображения:

class HorizontalFlip(CifarTransform):
    def _image_transform(self, image: Tensor, factor: float) -> Tensor:
        return F.hflip(image)


class VerticalFlip(CifarTransform):
    def _image_transform(self, image: Tensor, factor: float) -> Tensor:
        return F.vflip(image)


class AdjustBrightness(CifarTransform):
    def _image_transform(self, image: Tensor, factor: float) -> Tensor:
        return F.adjust_brightness(image, factor)


class AdjustContrast(CifarTransform):
    def _image_transform(self, image: Tensor, factor: float) -> Tensor:
        return F.adjust_contrast(image, factor)

Такие аугментации являются семейством операторов, преобразующих один тип данных Image: Tensor и образуют цепочку $T_0 \rightarrow T_1 \rightarrow T_2 \rightarrow ...$ . Поэтому их проще объединить в одну цепочку с помощью оператора композиции IcoPipeline:

item_aug_flow = IcoPipeline(
    HorizontalFlip(p=0.5),
    AdjustBrightness(p=0.2, factor_max_delta=0.2),
    AdjustContrast(p=0.2, factor_max_delta=0.2),
)
item_aug_flow.describe()

Теперь интегрируем наши аугментации в общий пайплайн. Как вы уже понимаете, цепочку item_aug_flow мы можем добавить в участок пайплайна, в котором происходит работа с типом данных CifarItem:

train_data_pipeline = (
    train_source 
    | shuffle_indices 
    | batcher 
    | (fetch_item | item_aug_flow).stream().stream()
    | collate
)
train_data_pipeline.describe()

Мы используем ссылку на уже описанный фрагмент пайплайна аугментации item_aug_flow и добавляем его сразу после оператора загрузки элемента датасета fetch_item .

Многопроцессорная параллельная обработка данных

Для выполнения произвольного пайплайна в отдельном процессе в ICO реализован оператор MPAgent. Он является оператором композиции и использует вложенный оператор, выполняя его в воркере в отдельном процессе. При этом, будучи оператором, он повторяет сигнатуру вложенного в него оператора.

Также MPAgent отвечает за жизненный цикл воркера, который выполняется в отдельном процессе. Для его запуска используется метод spawn: родительский процесс запускает новый интерпретатор Python, поэтому воркер не наследует состояние интерпретатора родительского процесса так, как это происходит при fork.

Это означает, что окружение, необходимое для работы воркера, должно быть явно создано или передано в новый процесс. Поэтому вложенный оператор передается в MPAgent в виде фабрики.

Такой подход также удобен для операторов, содержащих состояние или ресурсы, которые должны инициализироваться непосредственно внутри процесса воркера.

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

class WorkerFlowFactory:
    name: str | None = None
    batch_size: int

    def __init__(
        self,
        dataset: CifarDataset,
        batch_size: int,
        name: str | None = None,
    ):
        self.name = name
        self.batch_size = batch_size
        self.dataset = dataset
    
    def __call__(self) -> IcoOperatorProtocol[Iterator[int], CifarBatch]:
        
        @operator()
        def fetch_item(idx: int) -> CifarItem:
            """
            Получить элемент датасета по индексу
            """
            return self.dataset[idx]

        # Пайплайн аугментации

        item_aug_flow = IcoPipeline(
            HorizontalFlip(p=0.5),
            AdjustBrightness(p=0.2, factor_max_delta=0.2),
            AdjustContrast(p=0.2, factor_max_delta=0.2),
        )

        worker_flow = (fetch_item | item_aug_flow).stream() | collate
        worker_flow.name = "Worker flow"
        return worker_flow

worker_flow_factory = WorkerFlowFactory(
    dataset=dataset,
    batch_size=batch_size,
)
worker_flow_factory().describe()

На входе в пайплайн воркера приходит батч индексов Iterator[int] и он повторяет уже сделанный нами пайплайн, возвращая CifarBatch.

Теперь создадим MPAgent и выведем его план.

worker_flow_factory = WorkerFlowFactory(
    dataset=dataset,
    batch_size=batch_size,
)
mp_agent = MPAgent(worker_flow_factory)
mp_agent.describe()

При отображении плана для MPAgent также отображается вложенный пайплайн, который выполняется в воркере. Это позволяет достичь полной прозрачности логики выполнения, мы видим какие данные передаются воркеру, что является результатом его работы и в чем она состоит.

При рендеринге агента используется факт наличия фабрики для вложенного оператора, это позволяет отобразить его план выполнения.

Параллельная подготовка батчей данных

Для ускорения подготовки данных в PyTorch Dataloader используется пул воркеров. Каждый воркер подготавливает свой батч данных и возвращает его в общую очередь для дальнейшей обработки. В ICO есть аналог такого функционала, но он реализован как универсальная модель. Аналог воркера уже был описан - это MPAgent с вложенным произвольным оператором (точнее его фабрикой).

Для организации параллельных вычислений мы будем использовать MPAgent в сочетании с IcoAsyncStream. Этот оператор является управляющей конструкцией, схожей с IcoStream, которую мы уже рассматривали и имеет сигнатуру $Iterator[I] \rightarrow Iterator[O]$ . Но в отличие от обычного стрима, он асинхронно применяет пул операторов к элементам последовательности. Координацию их выполнения и ожидание результатов IcoAsyncStream берет на себя, используя механизмы AsyncIO. При этом в нашем случае сами вычисления выполняются параллельно в отдельных процессах благодаря MPAgent.

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

def create_mp_agent() -> IcoOperator[Iterator[int], CifarBatch]:
       # Создаем агента с фабрикой вложенного пайплайна воркера 
       
        worker_flow_factory = WorkerFlowFactory(
            dataset=dataset,
            batch_size=batch_size,
        )
        return MPAgent(worker_flow_factory)

# Создаем оператор асинхронного стрима с указанием размера пула операторов

workers_pool = IcoAsyncStream(create_mp_agent, pool_size=num_workers)
workers_pool.describe()

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

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

Теперь часть подготовки данных полностью готова и мы можем увидеть ее целиком в виде одного графа вычислений.

Полный пайплайн параллельной подготовки данных

Объединим все наши шаги в один пайплайн:

train_data_pipeline = (
    train_source
    | shuffle_indices
    | batcher
    | workers_pool
)
train_data_pipeline.name = "Train Flow"
train_data_pipeline.describe()

Полный пайплайн подготовки данных для тренировки

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

Но перед тем, как перейти к описанию цикла тренировки, надо уделить немного внимания оптимизации.

Оптимизация для многопроцессорного взаимодействия

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

Для этой цели используется API Tensor.share_memory_() в следующих частях кода:

Перенос датасета в общую память перед использованием в MPAgent:

class CifarDataset:
    def share_memory_(self) -> None:
        self.images.share_memory_()
        self.labels.share_memory_()

dataset.share_memory_()

Перенос батча в общую память при коллейте в воркере:

@final
class CifarBatch:
    def share_memory_(self) -> None:
        self.images.share_memory_()
        self.labels.share_memory_()

@operator()
def collate(items: Iterator[CifarItem]) -> CifarBatch:
    batch = CifarBatch(items)
    batch.share_memory_()
    return batch

Также изначально используется источник данных с индексами, а не самими элементами датасета. Это позволяет передавать в процесс воркера не копии CifarItem, а только индексы элементов.

Воркер, получая индекс, получает элемент из датасета, все данные в котором находятся в общей памяти и поэтому используется одна копия данных на все воркеры.

CifarBatch перемещается в общую память при коллейте в воркере и передаваясь через очередь, сериализуется только ссылка на общую память, а не сами данные в тензорах.

Вот теперь все готово для создания цикла тренировки.

Цикл тренировки нейронной сети

Ключевым элементом в этом разделе является новый объект - контекст. Он позволяет построить процесс обучения, сохраняя и передавая состояние между операторами.

Контекст обучения

Опишем контекст как структуру данных, она будет включать:

модель для обучения
оптимизатор
функцию потерь
вспомогательные поля

@dataclass(slots=True)
class CifarTrainContext:
    model: nn.Module
    optimizer: torch.optim.Optimizer
    loss_fn: torch.nn.Module = torch.nn.CrossEntropyLoss()
    iter_num: int = 0
    total_loss: float = 0.0

Дальше нам потребуются операторы для работы с контекстом.

Шаг обучения и контекстные операторы

В ICO есть семейство контекстных операторов. В простом случае, это просто оператор с сигнатурой $C_0 \rightarrow C_1$ , взаимодействующий с контекстом. Но для обучения такого оператора недостаточно, ведь для процесса обучения требуется не только контекст, но и данные.

Поэтому оператор обучения имеет сигнатуру $I_0, C_0 \rightarrow C_1$ и имеет класс IcoContextOperator. Он принимает элемент данных и контекст и возвращает модифицированный контекст. Опишем такой оператор для шага обучения нейронной сети:

def train_step(batch: CifarBatch, context: CifarTrainContext) -> CifarTrainContext:
    #  1. сброс градиентов оптимайзера
    context.optimizer.zero_grad()

    #  2. прямой проход модели и получение предсказания
    outputs = context.model(batch.images)

    #  3. Вызов функции потерь и получение ошибки
    loss = context.loss_fn(outputs, batch.labels)
    
    # 4. Обратный проход модели и вычисление градиентов
    loss.backward()
   
    #  5. Шаг оптимизации и обновление весов модели
    context.optimizer.step()

    #  6. Дополнительно - сохранение значения ошибки и номера итерации
    context.total_loss = loss.item()
    context.iter_num += 1
    return context

Также добавим шаг логирования и шаг сохранения чекпоинта и объединим шаги в пайплайн.

def logging_step(context: CifarTrainContext) -> CifarTrainContext:
    if context.iter_num % 10 == 0:
        print(f"{context.iter_num=}, {context.total_loss=:.4f}")
    return context

def save_checkpoint_step(context: CifarTrainContext) -> CifarTrainContext:
    if context.iter_num % 10 == 0:
        print(f"Checkpointing model at {context.iter_num=}")
    return context


train_pipeline = IcoContextPipeline(train_step, logging_step, save_checkpoint_step)
train_pipeline.name = "Training Pipeline"
train_pipeline.describe()

В коде выше использован оператор IcoContextPipeline, это оператор композиции, схожий с IcoPipeline, но c сигнатурой $I, C_0 \rightarrow C_1$ . Его задача упростить описание семейства операторов с одинаковой сигнатурой: первый оператор в такой цепочке всегда контекстный с сигнатурой $I, C_0 \rightarrow C_1$ , а остальные операторы имеют сигнатуру $C_1 \rightarrow C_2$ .

Мы описали один шаг обучения, теперь используем этот шаг для описания целой эпохи обучения.

Эпоха обучения

Эпоха обучения в нашем варианте - это последовательное применение шага обучения к набору батчей данных. Такая управляющая конструкция реализована в операторе IcoEpoch с сигнатурой $C_0 \rightarrow C_1$ . Она принимает два параметра инициализации:

источник данных source с типом IcoSource и сигнатурой $() \rightarrow Iterator[I]$
контекстный оператор context_operator с типом IcoContextOperator и сигнатурой $I_0, C0 \rightarrow C_1$

и выполняет следующий код:

for item in self.source(None):
    context = self.context_operator(item, context)
return context

В коде видим, что происходит последовательное чтение источника данных, применение контекстного оператора к каждому элементу данных и возвращение измененного контекста. Это именно то, что и требуется для реализации эпохи обучения! Поэтому опишем нашу эпоху так:

from ico import IcoEpoch,

# Перевод модели в режим обучения

@operator()
def start_train(context: CifarTrainContext) -> CifarTrainContext:
    context.model.train()
    return context

# Оператор эпохи обучения
train_epoch = IcoEpoch(
    source=train_data_pipeline,
    context_operator=train_pipeline,
)
train_epoch.name = "Training Epoch"

# Пайплайн тренировки

train_flow = start_train | train_epoch
train_flow.name = "Train Flow"
train_flow.describe()

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

источник данных for each in …: он разворачивается в полное описание пайплайна подготовки данных, который мы подготовили в предыдущем разделе.
контекстный оператор …apply: он разворачивается в пайплайн с шагами обучения, применяемыми к каждому элементу данных в последовательности

В итоге, оператор IcoEpoch модифицирует и возвращает контекст, в котором содержится состояние модели, измененное этапом обучения одной эпохи.

Валидация: пайплайн данных

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

чтение набора данных,
прямой проход модели

Но есть и отличия:

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

Все эти отличия делают пайплайн валидации значительно “легче”, но иногда на практике переиспользуется код обучения, что приводит к усложнению и избыточности в реализации валидации.

Начнем с описания пайплайна подготовки данных.


# Количество элементов датасета, отданных на валидацию.

val_split_size = int(dataset_size * val_split_ratio)

# Источник данных - набор индексов.
# Используем вторую часть индексов для валидации.

def get_val_indices() -> Iterable[int]:
    return all_indices[-val_split_size:]  

val_source = IcoSource[int](get_val_indices, name="CIFAR10 evaluation indices")

# Пайплайн данных для валидации

val_data_flow = val_source | fetch_item.stream()
val_data_flow.name = "Validation Flow"
val_data_flow.describe()

Как видно, пайплайн данных для валидации стал сильно проще по сравнению с тренировочным вариантом - это просто чтение CifarItem по индексу из датасета!

Валидация: контекст

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

Контекст валидации сохраняет ссылку на контекст обучения для следующей эпохи обучения и содержит поля для подсчета метрик.

@dataclass(slots=True)
class CifarEvalContext:
    train_context: CifarTrainContext
    accuracy: float = 0.0
    total_samples: int = 0

Валидация: шаги

Контекстным оператором в шаге валидации будет подсчет метрики accuracy. Этот оператор тоже имеет сигнатуру $I_0, C_0 \rightarrow C_1$ , т.к. ему требуются данные и контекст с моделью для подсчета метрики.

def calculate_accuracy(item: CifarItem, context: CifarEvalContext) -> CifarEvalContext:
    with torch.no_grad():
        outputs = context.train_context.model(item.image.unsqueeze(0))  # (1, 10)

        predicted = torch.argmax(outputs, dim=1)
        correct = predicted.item() == item.label
        context.total_samples += 1
        context.accuracy += float(correct)

    return context

Реализуем остальные шаги:

# Оператор начала валидации, который будет создавать контекст

@operator()
def start_eval(context: CifarTrainContext) -> CifarEvalContext:
    context.model.eval()
    return CifarEvalContext(train_context=context)

# Оператор логирования метрики

@operator()
def log_accuracy(context: CifarEvalContext) -> CifarEvalContext:
    accuracy = (
        context.accuracy / context.total_samples if context.total_samples > 0 else 0.0
    )

    validation_accuracy = accuracy * 100
    print(f"{validation_accuracy =:.2f}% ({context.accuracy}/{context.total_samples})")
    return context

# Оператор завершения валидации и получения контекста обучения

@operator()
def end_eval(context: CifarEvalContext) -> CifarTrainContext:
    """End evaluation phase and return to training context."""
    return context.train_context

Валидация: эпоха

Теперь у нас есть все необходимое для создания эпохи валидации.

# Описание эпохи валидации

val_epoch = IcoEpoch(
    source=val_data_flow,
    context_operator=IcoContextOperator(calculate_accuracy),
    name="Validation Epoch",
)

# Описание пайплайна валидации

val_flow = start_eval | val_epoch | log_accuracy | end_eval
val_flow.name = "Validation Flow"
val_flow.describe()

Как видно, пайплайн валидации по сравнению с тренировкой сильно упростился. Эпоха состоит из одного контекстного оператора с подсчетом метрики. Пайплайн начинается с создания контекста валидации и заканчивается извлечением контекста обучения.

Полный процесс обучения нейронной сети

Мы описали два графа: эпоха обучения и эпоха валидации. Они представлены двумя операторами, с сигнатурой $CifarTrainContext \rightarrow CifarTrainContext$ . Мы можем их соединить и получить итерацию из эпохи обучения и валидации.

epoch_flow = train_flow | val_flow
epoch_flow.name = "Train and Validation Epoch Flow"
epoch_flow.describe()

Итак, у нас есть итерация и все, что нам осталось сделать для обучения, это запустить цикл с заданным количеством итераций!

Для этой цели, в лучших традициях ICO, имеется еще один управляющий оператор IcoProcess с сигнатурой $С_0 \rightarrow C_1$ . Он применяет указанный оператор body заданное количество итераций:

for _ in range(self.num_iterations):
    context = self.body(context)
return context

Для описания процесса обучения нейронной сети, нам достаточно вложить построенный пайплайн в такой оператор:

train_process = IcoProcess(epoch_flow, num_iterations=num_epoch)
train_process.name = "CIFAR-10 Training Process"
train_process.describe()

Видим, что IcoProcess добавил уровень вложенности и включил в себя наш пайплайн, который состоит из эпохи обучения и эпохи валидации.

Первая и последняя строчка плана выполнения говорит, что граф требует на вход и возвращает CifarTrainContext. В текущем формате мы можем выполнить процесс обучения, создав контекст обучения и запустив оператор train_process(train_context). Но вспоминая математическое отступление из начала статьи, давайте приведем граф к финальной замкнутой форме с сигнатурой $() \rightarrow ()$ .

Финальный граф обучения

Для этого, в начало и конец графа нам надо добавить два оператора с сигнатурами $() \rightarrow I$ и $I \rightarrow ()$ соответственно, что и приведет его к форме $() \rightarrow ()$

@operator()
def create_train_context(_: None) -> CifarTrainContext:
    model = create_cifar10_resnet18()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    return CifarTrainContext(model=model, optimizer=optimizer)


@operator()
def finalize_train_context(context: CifarTrainContext) -> None:
    """Finalize training and report completion."""
    print("Training complete. Final model ready for test or deployment.")

cifar_flow = create_train_context | train_process | finalize_train_context
cifar_flow.name = "Complete CIFAR-10 Training and Validation Flow"
cifar_flow.describe()

По плану выполнения мы видим, что первый оператор ничего не принимает на вход, последний не возвращает. Такой граф мы можем выполнить просто вызвав cifar_flow(None).

Мы получили граф в замкнутой форме и тем самым выполнили поставленную цель - представили весь ML-пайплайн в виде единого графа на базе фреймворка ICO.

Заключение

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

При этом появилась формальная модель описания операторов через типизированные сигнатуры, в том числе для вычислений с состоянием: $(I_0, C_0) \rightarrow (I_1, C_1)$ . Это позволяет рассматривать такой пайплайн как более формальную математическую систему.

Представление вычислений в формате операторов вместе с принципом линейности позволило использовать композицию и свободно комбинировать операторы (в рамках допустимых типов, конечно). А представление управляющих конструкций в формате операторов добавило большей глубины и выразительности в графы ICO.

Мне интересно, какие еще задачи можно эффективно выразить в формате ICO и я был бы рад получить запросы от сообщества с предложением идей для экспериментов.

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

Очень естественным кажется добавление профилировщика производительности для произвольного графа ICO: уже есть описание графа вычислений и рендерер для плана выполнения. Остается только добавить сбор и отображение метрик производительности.

Еще была бы интересна дискуссия с рассмотрением текущих альтернатив и опыта читателя в построении вычислительных пайплайнов.

Исходный код ICO находится в этом репозитории на GitHub. Скрипт с кодом для этой статьи можно посмотреть и запустить в примерах.

В данной статье я не затрагивал тему среды выполнения и возможностей по управлению процессами и мониторингом прогресса, но это уже реализовано и могло бы быть темой следующей статьи.

Спасибо всем, кто дочитал до конца, и спасибо за терпение!

Превратил весь ML-пайплайн в единый вычислительный граф и ни о чём не жалею