Фреймворк ICO собирает пайплайн машинного обучения в единый граф без «швов»

Автор фреймворка ICO разобрал корневую проблему PyTorch-пайплайнов: подготовка данных живёт в Dataset и DataLoader, обучение живёт в ручном цикле, а между ними нет общей структуры. Данные передаются словарями со строковыми ключами, типы нигде не зафиксированы, параллелизм привязан к одному компоненту. ICO предлагает другой подход: каждый шаг становится оператором с явной сигнатурой «вход/выход», операторы соединяются в линейную цепочку, а управляющие конструкции (циклы, ветвления) тоже являются операторами.

Какую задачу решаем?

Если вы пишете пайплайн машинного обучения на PyTorch, скорее всего ваш код разбросан по нескольким местам:

• Загрузка данных и аугментации спрятаны внутри Dataset и collate_fn
• Прямой проход, функция потерь, обратное распространение ошибки (вычисление градиентов, чтобы модель «понимала», куда двигать свои параметры) и обновление весов написаны в отдельном цикле
• Структура передаваемых данных определяется «по договорённости» и нигде формально не описана

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

Что понадобится

• Python 3.10+ и установленный PyTorch (torchvision для работы с CIFAR-10)
• Фреймворк ICO (устанавливается через pip, репозиторий автора открыт)
• Базовое знакомство с Dataset и DataLoader в PyTorch
• Около 30 минут, чтобы пройти все шаги и запустить обучение классификатора на CIFAR-10

Пошаговая инструкция

1. Определите базовые структуры данных. Создайте дата-класс для одного элемента датасета (изображение и метка класса) и класс-обёртку, которая загружает весь CIFAR-10 в память. В ICO структура данных задаётся явно: никаких словарей с произвольными строковыми ключами.

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))]
        )
        self.labels = torch.tensor(
            [dataset[i][1] for i in range(len(dataset))]
        )

1. Оберните источник данных в оператор IcoSource. В ICO даже источник данных это оператор. Его сигнатура: он ничего не принимает на вход и возвращает итератор по элементам. На противоположном конце графа стоит IcoSink, оператор, который принимает данные и ничего не возвращает. Вместе они замыкают граф.

2. Опишите каждый шаг обработки как оператор. Аугментация, прямой проход модели (forward), вычисление функции потерь, обратное распространение ошибки (backward), обновление параметров оптимизатором, каждый из этих шагов становится отдельным оператором с явной сигнатурой: тип входных данных и тип выходных данных.

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

4. Добавьте контекст Context для состояния. Если оператору нужно запомнить что-то между шагами (текущая эпоха, накопленные метрики, состояние оптимизатора), он читает и изменяет объект Context. Оператор одновременно может преобразовывать данные и менять контекст, это позволяет выражать вычисления с состоянием в том же графе.

5. Запустите граф. Итоговая конструкция от IcoSource до IcoSink это замкнутый граф с сигнатурой «ничего на входе, ничего на выходе», самодостаточный оператор, выполняющий полный цикл обучения.

Кому это пригодится и что делать прямо сейчас?

ML-инженеру и дата-сайентисту. Если ваш пайплайн машинного обучения на PyTorch разросся до состояния, когда правка в аугментации ломает цикл обучения, попробуйте переписать один небольшой проект через ICO. Начните с CIFAR-10 по описанной схеме, это займёт полчаса и покажет, подходит ли вам такая модель мышления.

Автору технического контента на Дзене. Подход «всё есть оператор» хорошо ложится на обучающие статьи: каждый шаг изолирован, его можно объяснить отдельно. Если вы пишете туториалы по машинному обучению, формат графа операторов делает структуру статьи нагляднее.

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

Главный вывод: если вы ловите себя на мысли, что при каждом новом проекте заново пишете один и тот же хаотичный цикл обучения, попробуйте посмотреть на пайплайн как на граф операторов. Даже если ICO вам не подойдёт, сам принцип «каждый шаг это оператор с явным типом на входе и выходе» дисциплинирует код лучше любого код-ревью.