Исчерпывающее руководство по дообучению моделей 3D-генерации

TLDR; В этом блоге мы предоставляем скрипты и контейнеры для дообучения TRELLIS.2, популярной 3D-модели, на ваших собственных 3D-активах. Их можно запускать как локально, так и в Hugging Face Jobs. Если вы заинтересованы в дообучении, вы можете использовать наши скрипты с минимальным количеством аргументов (при условии, что вы загрузили свой набор данных в Hugging Face Hub) или поэкспериментировать с ними самостоятельно с помощью этого репозитория. Если вы хотите узнать больше, мы подробно разберем, как работает этот скрипт. Поехали!

TLDR; В этом блоге мы предоставляем скрипты и контейнеры для дообучения TRELLIS.2, популярной 3D-модели, на ваших собственных 3D-активах. Их можно запускать как локально, так и в Hugging Face Jobs. Если вы заинтересованы в дообучении, вы можете использовать наши скрипты с минимальным количеством аргументов (при условии, что вы загрузили свой набор данных в Hugging Face Hub) или поэкспериментировать с ними самостоятельно с помощью этого репозитория. Если вы хотите узнать больше, мы подробно разберем, как работает этот скрипт. Поехали!

Введение

3D-генерация находит множество применений в робототехнике, симуляции, медиа и индустрии развлечений. В этих областях сложно заставить сгенерированные объекты подчиняться строгим правилам и стандартизировать представления моделей и данных. Существует множество форматов файлов, таких как GLB и PLY. Также существует множество представлений, таких как нейронные поля излучения (NeRFs) и 3D Gaussian splatting (3DGS). У каждого из них есть свои преимущества и недостатки.

В TRELLIS было представлено общее внутреннее представление, называемое унифицированным латентным представлением. Его декодеры преобразуют это представление в 3D-модели в требуемом выходном формате. Таким образом, вам не нужно использовать разные модели для каждого формата представления.

Затем та же команда выпустила TRELLIS.2, которая улучшает представление 3D-структуры в модели. Но ни одна из моделей не предлагает простого способа адаптации к конкретной задаче.

Что мы хотим улучшить

Модель общего назначения, такая как TRELLIS.2, может допускать ошибки в деталях, которые важны для конкретного приложения. Например, пользователю может потребоваться точная геометрия для моделей мебели. Без дообучения под эту задачу предобученная модель может создавать недостаточно точную геометрию. Для нашего эксперимента мы используем подмножество GSO, содержащее 188 объектов обуви. В наших первоначальных тестах предобученная модель обычно создавала обувь, которая выглядела правильно, но подошвы были некорректными. Мы проверяем, может ли LoRA сделать подошвы более точными.

Наша цель — использовать набор данных Google Scanned Objects (GSO) для улучшения производительности TRELLIS.2 на конкретной категории объектов. Мы объясним, какие компоненты конвейера следует дообучать, а какие оставить без изменений во время обучения. Конвейер разделяет структуру, детальную геометрию и материалы. Таким образом, мы будем дообучать только те компоненты, которые важны для нашей задачи.

Выбор того, что заморозить, так же важен, как и выбор того, что дообучать.

Выбор того, что заморозить, так же важен, как и выбор того, что дообучать.

Что мы будем использовать?

LoRA

LoRA (Low-Rank Adaptation) дообучает модель путем тренировки небольших матриц адаптеров, сохраняя при этом исходные веса замороженными. Это уменьшает количество обучаемых параметров и объем памяти, необходимой для их градиентов и состояния оптимизатора. Мы используем библиотеку PEFT от Hugging Face для добавления и обучения этих адаптеров. Для получения дополнительной информации прочитайте наше руководство по LoRA.

Hugging Face Jobs

Hugging Face Jobs запускает контейнеризированные рабочие нагрузки на облачном оборудовании. Мы используем его для предварительной обработки и обучения. Чтобы использовать собственный GPU, ознакомьтесь с локальными командами Docker ниже. Прежде чем начать, убедитесь, что у вас есть:

• HF_TOKEN, который можно получить здесь
• Ваша учетная запись Hugging Face должна иметь доступ к DINOv3, так как DINOv3 является моделью с ограниченным доступом
• Если вы используете собственный набор данных, загрузите его меши в Hub перед началом.

Как работает TRELLIS.2

Прежде чем мы приступим к обучению, давайте посмотрим, как TRELLIS.2 генерирует 3D-объект и куда мы добавим LoRA-адаптеры.

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

Кодирование входного изображения

Сначала BiRefNet удаляет фон с входного изображения. Затем DINOv3-L кодирует изображение в токены признаков. Модели генерации используют эти токены для управления 3D-объектом.

Генерация разреженной структуры (Sparse Structure Flow)

Диффузионный трансформер (DiT) разреженной структуры прогнозирует грубую форму объекта. Это диффузионный трансформер, который создает небольшой латентный код размером 16³ × 8. Затем декодер вариационного автокодировщика (VAE) преобразует этот латентный код в сетку заполнения (occupancy grid) размером 64 × 64 × 64. Каждый воксель в сетке показывает, является ли это место частью объекта. Воксель — это пиксель в трех измерениях. Мы добавляем LoRA-адаптеры к этой модели, потому что ее выходные данные направляют этап геометрии.

O-Voxel и структурированные латентные представления

O-Voxel — это 3D-формат TRELLIS.2. Это разреженная воксельная сетка. Она сохраняет только воксели на поверхности объекта. Каждый воксель хранит локальную геометрию, такую как вершина меша, и атрибуты материала. Поскольку формат небольшой, процессор преобразует меш в O-Voxel за секунды. GPU преобразует O-Voxel обратно в меш за секунды.

Структурированные разреженные латентные представления поступают из TRELLIS. Они хранят 3D-данные в разреженных тензорах, что ускоряет вычисления. Модель использует только около 10% вокселей в сетке, поэтому она не тратит вычислительные ресурсы на пустые воксели. Таким образом, DiT формы и DiT текстуры являются разреженными трансформерами. DiT разреженной структуры — это плотный трансформер.

Генерация геометрии (Shape Flow)

DiT формы использует сетку заполнения и токены признаков DINOv3 для создания латентного представления. Затем другой декодер VAE преобразует этот латентный код в геометрию O-Voxel с разрешением 512 × 512 × 512.

Мы добавляем LoRA-адаптеры к DiT на этом этапе, потому что он контролирует форму объекта.

Генерация материалов + создание меша

На последнем этапе DiT текстуры использует латентное представление формы и признаки DINOv3 в качестве входных данных. Он определяет материалы и текстуры поверхности геометрии. Затем декодер VAE преобразует латентный код в значения базового цвета, металличности, шероховатости и альфа-канала для каждого вокселя. Мы не добавляем LoRA-адаптеры к этому этапу. В нашем наборе данных есть карты диффузного цвета, но в нем отсутствуют отдельные карты материалов, необходимые для обучения на этом этапе. Кроме того, наша проблема заключается в форме обуви. Базовая модель уже создает хорошие цвета и текстуры.

Наконец, конвейер объединяет атрибуты геометрии и материала в текстурированный меш и экспортирует его как файл GLB.

Планирование

Какие компоненты дообучать

Мы добавляем LoRA-адаптеры только к DiT разреженной структуры и DiT формы. Эти DiT контролируют форму и геометрию объекта. Мы замораживаем DiT текстуры, чтобы сохранить его текущие возможности.

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

Если мы обучим DiT потока текстур на этих картах, он сможет научиться отображать освещение как свойство поверхности. Он также может забыть то, что выучил ранее — проблема, называемая катастрофическим забыванием.

Зачем использовать LoRA

Наше подмножество содержит 188 моделей обуви из 1030 объектов в GSO. Для такого небольшого набора данных LoRA — лучший выбор, чем полная донастройка (fine-tuning). Полная донастройка обходится гораздо дороже. Она также может перезаписать то, что модель уже делает хорошо (ее выученные априорные знания). Наша конфигурация обучает в общей сложности около 9,3 миллиона параметров адаптера для обеих целевых DiT. Мы используем адаптеры ранга 8 в слоях внимания и полносвязных (MLP) слоях DiT разреженной структуры. Мы используем адаптеры ранга 16 в слоях внимания DiT формы. Ранг определяет размер матриц адаптера. Более высокий ранг позволяет адаптеру обучаться лучше.

Что вы будете запускать

Настройка 3D-модели может потребовать множества зависимостей и этапов установки. Мы предоставляем два образа Docker (которые вы можете запустить на Jobs или на своей собственной машине) с необходимым программным обеспечением. Один образ предназначен для предварительной обработки, а другой — для донастройки.

Чтобы обучить ваши адаптеры LoRA и загрузить их в Hub:

• Предоставьте контейнерам доступ к вашим данным.
• Установите параметры для вашей задачи.
• Запустите контейнер предварительной обработки.
• Запустите контейнер донастройки.

Затем мы объясним скрипт донастройки, его настройки и то, как каждая настройка влияет на обучение.

Обзор эксперимента

Мы пройдем через три части эксперимента:

• Подготовка данных. Преобразование сеток GLB или OBJ в латентные представления, которые использует TRELLIS.2 для обучения, и рендеринг входных изображений.
• Обучение адаптеров. Запуск одного задания для каждой DiT с подготовленным набором данных и конфигурацией LoRA.
• Объяснение кода. Изучение того, как контейнеры загружают данные, добавляют адаптеры и обучают модели.

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

Пошаговое руководство

Запуск в Hugging Face Jobs

Предварительная обработка

Контейнер предварительной обработки использует Blender для рендеринга 16 видов каждого объекта. Обучение LoRA будет использовать эти виды в качестве обусловливающих изображений. Blender работает без дисплея через утилиту xvfb-run. Контейнер также обрабатывает каждый объект через конвейер предварительной обработки. Пример промежуточного вывода см. в наборе данных обуви здесь. Для наших 188 объектов полный прогон занял около 1 часа на задании a100-large. Компактная версия конвейера предварительной обработки также доступна в этом сопутствующем репозитории.

Для Hugging Face Jobs мы передаем следующие аргументы:

Этот контейнер принимает еще несколько аргументов (независимо от того, запущен ли он в Hugging Face Jobs или через Docker), которые вы можете найти в документации здесь. Шаги, предпринятые в конвейере предварительной обработки, задокументированы в этом разделе ниже.

Донастройка LoRA

Контейнер обучения загружает подготовленный набор данных и предустановку или файл конфигурации. Он записывает метрики обучения с помощью Trackio. Используйте панель управления Trackio для мониторинга запуска.

Мы используем «ss» и «shape» для обозначения двух режимов, которые у нас есть для обучения LoRA. Режим «ss» будет донастраивать только DiT потока разреженной структуры, тогда как режим «shape» будет донастраивать DiT потока формы. По сути, они просто ссылаются на разные файлы config.json, используемые для каждого этапа. Более подробная информация о процессе доступна в этом разделе ниже.

Если вы установите --push-to, контейнер загрузит контрольные точки адаптера в Hub. Контейнер принимает параметры, которые также кратко задокументированы здесь. Компактный скрипт также доступен в том же репозитории.

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

Запуск локально с Docker

Команды и флаги такие же, как для Hugging Face Jobs. Отличаются только опции Docker.

Если вы хотите использовать наш уже предварительно обработанный набор данных локально на выбранной вами машине, используйте следующие команды:

Наши обученные адаптеры LoRA находятся в этом репозитории: suvadityamuk/trellis-2-finetune-gso

Внутри конвейера обучения

Предварительная обработка в латентные представления

Мы используем замороженные энкодеры TRELLIS.2 для преобразования входных сеток в латентные представления для обучения. Мы сохраняем эти латентные представления, чтобы последующим запускам не приходилось снова кодировать сетки.

Мы также рендерим 16 видов каждого объекта в Blender. Во время обучения мы выбираем один случайный вид для каждого образца в качестве обусловливающего изображения. Мы оставляем 9 из 188 объектов для оценки и обучаемся на остальных 179. Эти отложенные объекты показывают, работают ли адаптеры на объектах, не виденных моделью во время обучения.

Некоторые из видов, которые Blender отрендерил для одного объекта

Некоторые из видов, которые Blender отрендерил для одного объекта

Первый фрагмент настраивает данные. FLOW выбирает DiT потока для обучения (ss или shape), а STEPS устанавливает количество шагов оптимизатора. Размер пакета 4 с 4 шагами накопления градиента дает эффективный размер пакета 16. Затем фрагмент загружает файл конфигурации LoRA для выбранной DiT потока, создает набор данных в формате TRELLIS.2, удаляет отложенные объекты и создает загрузчик данных для донастройки.

Подготовка адаптеров LoRA

Мы используем библиотеку peft для вставки адаптеров LoRA в выбранные слои. У TRELLIS.2 есть собственный прямой проход (forward pass), поэтому стандартная обертка PeftModel() не работает. Вместо этого мы вызываем inject_adapter_in_model(). Эта функция добавляет адаптеры непосредственно в модель и сохраняет ее собственный интерфейс.

Мы храним обучаемые веса адаптера в 32-битном формате с плавающей запятой (fp32) для численной стабильности и используем автокастинг bfloat16 (bf16) для прямого прохода. Автокастинг позволяет PyTorch использовать более низкую точность для некоторых операций, что снижает использование памяти. Автокастинг не меняет способ хранения весов замороженного бэкенда.

Обучение

Важно отметить, что наша целевая функция здесь — это старое доброе сопоставление выпрямленного потока (как в популярных моделях генерации изображений, таких как Stable Diffusion 3, Flux 2 и других). TRELLIS.2 использовала ту же цель при предварительном обучении, поэтому это хороший выбор для донастройки. Чтобы узнать теорию, прочитайте этот блог о сопоставлении потоков (flow matching).

Мы обновляем только параметры LoRA в DiT разреженной структуры и формы. Все исходные веса модели остаются замороженными.

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

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

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

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

Затем мы выполняем следующие шаги в цикле обучения:

• Выбор временного шага t для каждого образца (используя logitNormal для DiT потока разреженной структуры и uniform для DiT потока формы на основе конфигурации TRELLIS.2).
• Генерация случайного шума, соответствующего форме данных, и изменение формы t для широковещательной рассылки.
• Смешивайте чистые данные с шумом на основе t (где t = 0 — это данные с минимальным шумом, а t = 1 — чистый шум).
• Прогнозируйте скорость при t × 1000, чтобы соответствовать условиям предварительного обучения.
• Рассчитывайте функцию потерь как среднеквадратичную ошибку между прогнозируемой скоростью и целевой скоростью (первой производной x_t).

Цикл обучения вычисляет функцию потерь и обновляет веса адаптера. Накопление градиентов суммирует градиенты из 4 небольших пакетов перед каждым шагом оптимизатора. Контрольные точки градиента (gradient checkpointing) экономят память GPU за счет повторного вычисления некоторых промежуточных значений во время обратного распространения ошибки. Внешний цикл выполняется заданное количество шагов оптимизатора, в то время как внутренний цикл накапливает градиенты по набору пакетов.

Для каждого пакета DINOv3 кодирует изображение условий. Для исключения классификатора (CFG) цикл устанавливает условие в ноль для случайной доли (p_uncond) образцов. Затем он вычисляет функцию потерь flow-matching с использованием bf16 autocasting. После внутреннего цикла он обрезает градиенты, чтобы предотвратить слишком большие обновления, и выполняет шаг оптимизатора. Он записывает значение функции потерь каждые 10 шагов и сохраняет контрольную точку адаптера каждые 500 шагов.

Результаты

Мы обучили два адаптера: один для DiT с разреженной структурой и один для DiT формы. Файлы конфигурации LoRA в образе контейнера содержат полные настройки обучения: ss_flow_lora_gso.json для DiT с разреженной структурой и slat_flow_img2shape_lora_gso.json для DiT формы. Каждый запуск использовал один A100 80 ГБ в Hugging Face Jobs. Запуск для разреженной структуры занял 7,2 часа, а запуск для формы — 3,6 часа.

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

Функция потерь flow-matching варьируется в зависимости от выбранных объектов, шума и временных шагов. Чтобы выяснить, улучшает ли дообучение геометрию, мы должны также сравнить результаты на отложенных объектах.

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

Мы также используем три метрики для оценки геометрии на тестовом наборе:

В сравнении девяти объектов на диаграмме базовая модель имеет средний F-score 0,276 по сравнению с 0,803 для модели с адаптерами для обуви. Расстояние Чамфера падает с 0,174 до 0,033, а нормальная согласованность возрастает с 0,530 до 0,804. С обоими адаптерами F-score лучше на всех девяти отложенных объектах.

Статистику всех запусков см. в этом Trackio space.

Будущие направления

Фух, ну вот и всё! Мы дообучили TRELLIS.2 для генерации обуви, добавив адаптеры LoRA к двум компонентам модели, которые управляют геометрией. Мы подготовили данные один раз, обучили адаптеры в отдельных заданиях и оценили результаты на отложенных объектах.

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

Далее мы можем изучить, помогут ли те же варианты другим моделям 3D-генерации, таким как Pixal3D и UniLat3D.

Модели Gaussian-splat и генерации мира, такие как HunyuanWorld 2.0 и NVIDIA Lyra 2.0, также являются возможными следующими шагами. Их представления и конвейеры обучения могут потребовать других методов.

Спасибо, что были с нами, увидимся в следующий раз! 👋🤗