Единый граф вычислений: как фреймворк ICO превратил ML-пайплайн в самодостаточный поток
Вместо фрагментированных скриптов и разрозненных компонентов современные исследователи экспериментируют с объединением всего процесса обучения нейронных сетей в один вычислительный граф. На примере платформы ICO и задачи классификации изображений CIFAR-10 показано, как такой подход обеспечивает явную структуру данных, гибкую аугментацию и многопроцессорную обработку без потери читаемости кода.
# Единый граф вычислений: как фреймворк ICO превратил ML-пайплайн в самодостаточный поток
В экосистеме машинного обучения традиционный подход к обучению нейронных сетей часто воспринимается как набор разрозненных этапов: загрузка данных, их подготовка, циклы обучения и валидация. Однако, как показывает опыт разработчика Raen1, этот пайплайн может быть представлен как единый, самодостаточный вычислительный граф. Такой подход, реализованный в фреймворке ICO, позволяет устранить неявную структуру данных и жесткую привязку логики к конкретным библиотекам.
Эксперимент проводился на классической задаче классификации изображений из набора CIFAR-10. Целью было создание архитектуры, где любой шаг — от индексации файла до вычисления функции потерь — является явным оператором, соединенным в линейный однонаправленный поток.
Архитектура единого графа: от оператора к замыканию
Основой подхода ICO является представление любой вычислительной операции в виде оператора. Каждый оператор имеет строго определенную сигнатуру: он принимает входные данные определенного типа и возвращает выходные данных, предопределяя структуру потока.
Граф вычислений строится как цепочка таких операторов. Выход одного служит входом для следующего. Особое внимание уделяется управляющим конструкциям, таким как циклы, которые также абстрагируются в виде операторов. Для сохранения состояния между этими операциями используется контекст, позволяя графу моделировать сложные процессы, включая обучение с обратным распространением ошибки.
Интересно, что вся система может быть понята через призму теории категорий или простой логики потоков. Источник данных рассматривается как оператор, не требующий ввода (аналог $\emptyset$), а финальный оператор, завершающий вычисления, выступает как приемник. Таким образом, весь пайплайн обучения можно описать как единый оператор с сигнатурой «пустота на входе, результат на выходе». Это создает замкнутый, самодостаточный граф, где каждая часть имеет четкое назначение и место в общей цепи.
Промышленная подготовка данных и аугментация
Реализация пайплайна началась с определения базовых структур данных. Для работы с датасетом CIFAR-10 были введены дата-классы, описывающие отдельный элемент (изображение и метку класса) и весь датасет в целом. Важно, что данные загружаются полностью в память, что упрощает работу с тензорами PyTorch.
Ключевой особенностью является использование индексов вместо прямого обращения к данным. Это позволяет разделить потоки на тренировочные и валидационные, а затем применить к ним операции случайного перемешивания и группировки в батчи. Операторы в ICO типизированы, что позволяет статическим анализаторам, таким как mypy, проверять корректность соответствия типов данных еще до выполнения кода.
Особую роль играет оператор stream. Он необходим для работы с вложенными итераторами. Например, при формировании батчей данные находятся в структуре Iterable[Iterable[int]] (список списков индексов). Оператор stream позволяет применять вложенные вычисления: внешний слой обрабатывает батчи, а внутренний — элементы внутри них. Синтаксический сахар .stream().stream() поднимает оператор загрузки элемента на необходимый уровень вложенности, обеспечивая совместимость сигнатур.
Аугментация данных также представлена как цепочка операторов. Реализован базовый класс трансформации, определяющий вероятность применения и диапазон изменений (например, сдвиг яркости или контраста). Специфические операторы, такие как горизонтальное или вертикальное отражение изображения, наследуют эту логику. Эти операторы объединяются в подпайплайн, который встраивается в основную цепочку сразу после загрузки элемента датасета.
Формирование батча (collate function) выполняется специализированным оператором, который собирает список элементов в тензоры изображений и меток. В результате выходом становится итератор объектов CifarBatch, готовый к подаче в модель.
Многопроцессорная обработка и прозрачность
Традиционные фреймворки, такие как PyTorch, используют DataLoader с фиксированным механизмом многопроцессорной подготовки данных, управляемым через API. ICO предлагает более гибкое решение благодаря оператору MPAgent.
MPAgent позволяет выполнять произвольный пайплайн в отдельном процессе (воркере). Это достигается путем передачи фабрики операторов в новый интерпретатор Python. Такой подход решает проблему наследования состояния, характерную для fork, и позволяет изолировать ресурсы, необходимые для работы воркера.
Пайплайн воркера, включающий загрузку данных, аугментацию и формирование батчей, создается внутри фабрики и возвращается как готовый оператор. MPAgent управляет жизненным циклом этого воркера, обеспечивая параллельную обработку потоков данных без потери прозрачности логики.
Одним из главных преимуществ такого подхода является возможность визуализации и анализа плана выполнения. Утилита describe() выводит детализированную схему потока с использованием библиотеки Rich, показывая классы операторов, их функции, параметры и типы данных. Это делает отладку и понимание архитектуры пайплайна гораздо более доступными, чем в случае с «черным ящиком» стандартных библиотек.
Такая структура не только упрощает внедрение параллелизма, но и делает код более модульным и понятным, превращая сложную процедуру обучения в набор составных частей, где каждый элемент подконтролен и понятен разработчику.