Как мы собираем данные для обучения Kandinsky31.05.2024 11:15

Всем привет! Наша команда в Sber AI занимается генеративными моделями и сегодня мы расскажем про очень важный этап разработки моделей для генерации фотореалистичных изображений и видео — процесс сбора и фильтрации данных. Про этот этап редко подробно рассказывают разработчики и исследователи таких известных генеративных моделей как DALL-E 3, Stable Diffusion, MidJourney или SORA. Генеративные модели уже многих впечатлили своими возможностями создавать максимально реалистичные изображения и видеоролики, а качественные данные — далеко не последняя причина, по которой такого качества генераций удалось достичь.

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

Поэтому в этой статье мы постараемся частично закрыть эти пробелы и рассказать как мы в Sber AI собираем, фильтруем и храним данные для обучения моделей Kandinsky и Kandinsky Video.

Разделим повествование о нашей работе с данными на следующие блоки:

1. Обработка и фильтрация датасетов пар текст‑изображение

2. Обработка и фильтрация датасетов пар текст‑видео

3. DPF — инструмент для обработки мультимодальных данных

Сбор, фильтрация и обработка пар текст-изображение

Для генерации изображений по тексту используют разные архитектуры и подходы, однако большинство объединяет тот факт, что для обучения text-to-image моделей используют огромные датасеты пар текст-изображение, размеры которых могут достигать миллиардов пар. Более того, качество обученной на этих данных модели зависит далеко не только от размера датасета, но и от качества данных. Что такое качественные данные — это важная тема для обсуждения, которую мы затронем в данной статье. 

Хранение датасетов

Датасеты для обучения text-to-image моделей составляют сотни миллионов и миллиардов пар текст-изображение. Даже предположив, что каждая картинка весит 100Кб, что далеко не так, потому что мы хотим работать с HD-изображениями и не прибегать к методам сжатия, то общий объём датасета без учёта текстов будет около 100Tb. Такой датасет довольно дорого хранить локально. Стандартным решением этой проблемы является хранение данных на плоском объектном хранилище, например Amazon S3, и использование WebDataset для чтения данных напрямую с объектного хранилища во время обучения.

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

PMC4991227_00003.json
PMC4991227_00003.jpg
PMC4537884_00002.json
PMC4537884_00002.jpg
PMC4323233_00003.json
PMC4323233_00003.jpg
PMC5429906_00004.json
PMC5429906_00004.jpg

Каждая пара .json и .jpg образует один обучающий пример. В .json файле хранится описание изображения и различная метаинформация. В одном tar архиве мы храним по 1000 примеров.

Во время обучения подгрузчик даных webdataset параллельно считывает эти архивы, которые далее мы будем называть шардами (shards), и они уже распаковываются in-memory. Скорость чтения файлов с S3 обычно в разы медленнее, чем с NAS, но за счет параллельного чтения нескольких шардов и того факта, что в одном архиве по 1000 сэмплов, скорость загрузчика успевает за скоростью обучения модели. То есть, пока идёт forward и backward проходы у модели, загрузчик успевает подгрузить новый батч данных. 

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

Однако у формата webdataset есть довольно существенный недостаток — метаинформация хранится внутри tar-архива вместе со всеми бинарными данными и в случае, когда мы хотим получить доступ только к метаинформации (например, получить все описания изображений), то придётся скачивать архивы вместе со всеми изображениями. Такие случаи возникают часто, поэтому мы используем немного модифицированный формат — храним бинарные файлы (изображения и видео) в архивах, а описания и метаинформацию в CSV-таблицах рядом с архивом. Получается следующая структура:

000000.tar
000000.csv
000001.tar
000001.csv
000002.tar
000002.csv
…

Чтобы иметь возможность соотнести метаинформацию и бинарники, в CSV-файле храним имена изображений из архива. Такой подход усложняет поддержку консистентности данных, но зато даёт возможность очень быстро просматривать метаинформацию по всему датасету, так как достаточно будет считать только .csv-файлы.

Фильтрация и отбор данных

Выше мы упоминали хорошие и плохие данные и их влияние на качество модели. Но какие данные можно считать хорошими, а какие — плохими? В относительно недавних статьях Pixart-alpha и YandexART тоже подробно рассматривают этот вопрос, а также проводят эксперименты с влиянием разных данных на процесс обучения генеративных диффузионных text-to-image моделей.

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

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

• Многие изображения содержат водяные знаки, которые модель может запомнить при обучении и воспроизвести при генерации. Нужно этого избегать;

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

• Низкое качество и артефакты на изображении: наличие размытия, слишком темных или пересвеченных фото и т. д. Это свойство можно рассматривать как эстетичность изображения.

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

Без фильтрации по эстетичности

С фильтрацией по эстетичности

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

• Схожесть текста и изображения. Для решения данной задачи обычно используют CLIP-score и подобные подходы;

• Значение эстетичности изображения. Можно вручную разметить небольшое количество данных и обучить модель предсказывать эстетичность и общее качество фото;

• Вероятность водяного знака на изображении. Используется модель детекции водяных знаков (например, вот эта), которая возвращает вероятность от 0 до 1;

• Вероятность NSFW контента. Для этого тоже используется отдельная модель, обученная на специальной разметке. Например, модель от LAION для этой задачи;

• Площадь текста на изображении. Картинки с большим количеством текста можно находить по площади текста на фото и CLIP скорам между фото и промптами «image with text» и др;

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

Эти основные метрики мы используем для фильтрации нашего внутреннего датасета, пороги к фильтрам подбирали эмпирически. Все эти фильтры реализованы в DataProcessingFramework (DPF), документация к ним есть в репозитории на гитхабе. Вообще таких атрибутов и метрик много, если вам интересно узнать какие ещё атрибуты можно подсчитывать для фильтрации text-image датасета, можете посмотреть статью YandexART. 

Также на разных этапах обучения мы используем разные пороги для фильтрации. Это позволяет регулировать и подбирать нужную «чистоту» и эстетичность данных. Мы тренировали Kandinsky в несколько этапов:

1. Обучение на всем датасете в низком разрешении с минимальными фильтрами. На этом этапе убираем все изображения с водяными знаками и убираем изображения с большим количеством текста;

2. Дообучение на изображениях в более высоком разрешении и высоким показателем эстетичности.

По проведённым экспериментам такой подход показал лучшее качество модели. Качество мы замеряли с помощью side-by-side сравнения на фиксированной корзине промптов.

Синтетические описания

Большой проблемой в данных текст-изображение и текст-видео является частое несоответствие описания изображению (видео) или отсутствие подробностей. Это связано с тем, что описание к большинству изображений берется из атрибута alt у html-тэга, а в нем часто содержится не соответствующее описание, либо описания вовсе нет. При этом, такие изображения могут быть в хорошем качестве и иметь высокую оценку эстетичности, но мы их не сможем взять в обучение из-за неинформативного описания. Эту проблему также затрагивают в статье Pixart-alpha: в ней было предложено использовать модель LLaVA для генерации синтетических описаний, которые более информативны и точнее описывают изображение. Такой подход показывает отличные результаты, поэтому мы его реализовали, использовав помимо LLaVA и другие vision-language models (VLMs).

Также можно было бы использовать LLM для расширения существующего промпта, но тогда языковая модель будет додумывать объекты, которых нет на изображении. Мы сравнили эти подходы, используя набор изображений из интересующих нас доменов, и оценили получившиеся кэпшены side‑by‑side. Лучше всего себя показали VLM модели, такие как LLaVA.

В таблице ниже можно наглядно увидеть разницу между описаниями из интернета и описаниями, которые генерирует LLaVA-1.5 с промптом «Describe this image and its style in a very detailed manner».

Русский культурный код

Идея собирать и прокачивать у Kandinsky понимание российской культуры пришла после релиза версии 2.2. Ниже приведено несколько картинок по популярным сущностям культурного кода Kandinsky 2.2:

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

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

Чтобы было понимание, что искать, мы создали таблицу русских сущностей и накидали туда всё, что ассоциируется с российской культурой. Затем пересмотрели, разбили на классы и почистили. Выделились такие темы, как, например, еда (блины, селедка под шубой), достопримечательности (Казанский собор, Красная площадь), персонажи популярных мультфильмов и фильмов. На текущий момент наш список насчитывает порядка 8000 сущностей по 16 тематикам.

Изображения собирались из различных открытых источников, из доступных всем советских и российских фильмов, а также мультфильмов. Однако эти изображения не всегда имели хорошее качество и подходящее описание. Генерации Kandinsky 3.0, обученному на этих данных, по запросу «Чебурашка» показаны на рисунке ниже:

Как можно заметить, в целом у модели появилось понимание, как примерно выглядит Чебурашка, однако изображения выглядят пугающе и совсем не красиво. 

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

Результаты генерации Kandinsky 3.0, обученном на составленном датасете представлены ниже:

Также проведено side-by-side тестирование, в результате которого было установлено, что понимание культурного кода у Kandinsky 3.0 стало выше на 23% по сравнению с версией 2.2.

В итоге у нас получилось за ~2–3 месяца собрать почти 100к пар данных по российскому культурному коду, который помог улучшить качество генераций модели Kandinsky на этом домене.

Данный подход по ручному отбору заранее подготовленной выборки изображений с последующим их описанием был протестирован на домене изображений с текстом и также дал на нём хороший прирост качества. Пример генерации ниже:

Сбор, фильтрация и обработка пар текст-видео

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

Хранение

Датасеты с видео весят в разы больше, чем с изображениями, поэтому мы их тоже храним в объектном S3-совместимом хранилище. Используем структуру с tar-архивами и csv-файлами, так как она показывает неплохие значения по скорости чтения и нам также важна возможность фильтровать данные по атрибутам прямо во время обучения.

Препроцессинг

Для эффективного обучения генеративной видео модели крайне важно правильно предобработать датасет. В отличие от картинок, для видео этот этап является чуть ли не основным. Модель генерации видео обучается на коротких видео длиной от 2-х до 60-ти секунд. Желательно, чтобы действия в этих видео происходили в рамках одной сцены, то есть было логичное завершенное действие без резкого перехода к другой сцене. Для этого сначала длинные видео нарезаются на такие сцены с помощью утилиты PySceneDetect, а затем уже обрабатываются полученные короткие видео. 

Наш пайплайн предобработки сцен состоит из нескольких этапов:

1. Отбрасываем слишком короткие (< 2 секунд) и слишком маленькие (площадь кадра < 256*256 px или одна из сторон < 128 px) видео.

2. Отбрасываем видео с FPS < 23 — таких видео очень мало, но они затрудняют обучение, так как обучение должно проводиться на данных с одним FPS.

3. Масштабируем большие видео до 1080 px по минимальной стороне. Видео в 4K занимают много памяти, в то время как предобучение модели проводится в разрешениях 256, 512 и 1024. Хранить 4K видео в шардах не имеет смысла, так как это в разы замедлит скорость работы загрузчика.

4. Конвертируем все видео к 24 FPS.

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

6. Перемешиваем образцы в наборе, чтобы батч с большей вероятностью собирался из более разнородных данных.

Для ускорения пайплайна препроцессинга видео мы объединяем этапы 3–5, используя сложный фильтр в ffmpeg — это позволяет выполнить преобразования за одно чтение файла и значительно ускорить процесс.

Фильтрация и обработка

Фильтры для видео, по сути, включают в себя все фильтры для изображений, а также фильтры, учитывающие динамику видео:

• эстетичность;

• наличие водяного знака;

• наличие текста на видео.

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

• Оптический поток — величина, определяющая степень изменения объекта между кадрами на видео. Можно использовать различные алгоритмы для измерения оптического потока: Gunnar Farneback’s algorithm и модель RAFT;

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

Синтетические описания

Описания для видео в интернете еще хуже, чем для картинок — информативных описаний крайне мало. Поэтому мы генерируем синтетические описания для видео и учимся на них. 

Для генерации описаний к видео мы протестировали несколько подходов:  

1. Использовали vision-language models (VLMs) для картинок, применяя их на центральный кадр видео. Но тогда полученные описания не будут описывать происходящее на всём видеоролике и будут упускать множество деталей и информации;

2. Можно было бы описать несколько кадров, а затем использовать обычную LLM для «склейки» нескольких описаний в одно. В этом случае также возможны галлюцинации и несоответствие описания и видео;

3. Использовать VLM, поддерживающую работу с видео. Например, модели Video-LLaVA, LITA, pLLaVA.

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

DPF — фреймворк для обработки мультимодальных данных

У нас часто возникают задачи с фильтрацией и обработкой картиночных или видео данных, причем задачи не обязательно связаны с подготовкой данных для обучения text-to-image или text-to-video моделей. Мы проводим множество других экспериментов, которые требуют предварительной подготовки данных — фильтраций, перевода, генерации описаний к изображениям, создание инструкций и многое другое. Эти задачи могут выполнять разные сотрудники, поэтому хотелось бы иметь один репозиторий со всеми фильтрами и обработчиками. Для этого мы разработали библиотеку для обработки и фильтрации данных — DPF.

Фреймворк состоит из следующих основных сущностей:

• Processor — объект, инкапсулирующий в себя датасет и все методы для взаимодействия с данными. Датасет представляется в виде таблицы, а файлы в нем указываются путями;

• Filter — фильтр для обработки данных. Каждый фильтр просто добавляет один или несколько новых атрибутов данных. Каждый атрибут — отдельная колонка в таблице. Примерами фильтров могут быть фильтр эстетичности, который возвращает значение эстетичности изображения, и фильтр схожести текста и изображения на основе CLIP;

• Transforms — трансформации данных. Используются для изменения самих данных. Например, для изменения FPS видео или изменения размера изображений;

• Pipelines — последовательность трансформаций и фильтров с логированием и обработкой ошибок.

Во время разработки мы стремились добиться следующих целей:

1. Единый код для всех фильтров, которые применяются к датасетам. Часто возможны ситуации, когда разработчики из разных команд параллельно реализуют один и тот же фильтр, но реализуют его с отличиями. Эта разница в реализации влияет на подсчет атрибутов для фильтрации. Затем, когда наступит время поделиться датасетом с другой командой, то будет тяжело соотнести атрибут и фильтр, которым этот атрибут был посчитан. Поэтому важно иметь единый репозиторий для реализации фильтров для данных;

2. Уменьшение времени разработки новых фильтров. В DPF достаточно реализовать только логику взаимодействия фильтра с данными, а чтение данных, запись атрибутов в отдельную колонку и прочую общую логику для всех фильтров повторно реализовывать не нужно;

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

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

Давайте рассмотрим несколько примеров кода. Запустить фильтр генерации описаний можно в несколько строчек кода, достаточно считать датасет и указать настройки фильтра:

from DPF import DatasetReader, ShardsDatasetConfig
from DPF.filters.images.llava_captioning_filter import LLaVaCaptioningFilter

reader = DatasetReader()

# Создаем конфиг датасета
config = ShardsDatasetConfig.from_path_and_columns(
    "examples/example_dataset",
    image_name_col='image_name',
)
# считывает метаданные датасета
processor = reader.read_from_config(config, workers=16)

# создаем фильтр для создания описаний с помощью LLaVA
datafilter = LLaVaCaptioningFilter(
    workers=16, batch_size=16, device="cuda:0"
)
# запускаем фильтр
processor.apply_data_filter(datafilter) 

К датасету добавятся новые метаданные — сгенерированное описание. Посмотреть все метаданные можно, выведя атрибут df объекта processor: print(processor.df). Метаданные хранятся просто в pandas-датафрейме, что позволяет удобно с ними работать. Чтобы добавить или обновить метаданные в самом датасете, используйте метод update:

# Добавление новых метаданных в датасет
processor.update_columns([new_column_name], workers=16)

Можно запустить фильтр на нескольких GPU для более быстрой обработки датасета:

from DPF.filters.multigpu_filter import MultiGPUDataFilter

multigpufilter = MultiGPUDataFilter(
    ['cuda:0', 'cuda:1', 'cuda:2', 'cuda:3'],
    LLaVaCaptioningFilter,
    dict(
        pbar=True, workers=8,
        prompt='short', batch_size=16
    )
)
processor.apply_multi_gpu_data_filter(multigpufilter)

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

from DPF.configs import ShardsDatasetConfig
from DPF.dataset_reader import DatasetReader
from DPF.pipelines import FilterPipeline
from DPF.filters.images.info_filter import ImageInfoFilter
from DPF.filters.images.hash_filters import PHashFilter

reader = DatasetReader()
config = ShardsDatasetConfig.from_path_and_columns(
    "examples/example_dataset",
    image_name_col='image_name',
)
processor = reader.read_from_config(config, workers=4)

pipeline = FilterPipeline("pipeline_example")
pipeline.add_datafilter(
    ImageInfoFilter,
    {'workers': 4},
    processor_run_kwargs={'return_none_on_error': True},
)
pipeline.add_dataframe_filter(lambda df: df[df['is_correct']])
pipeline.add_datafilter(PHashFilter, {'workers': 4})
pipeline.add_deduplication(["image_phash_8"])
pipeline.add_shuffle()
pipeline.run(processor)

Заключение

Надеемся, что эта статья внесла больше прозрачности в процесс подготовки данных для text-to-image и text-to-video моделей. Мы планируем активно развивать фреймворк и добавлять в него новые фильтры, поэтому будем рады обратной связи и контрибьюту. Stay tuned!

Авторы

Сбор и обработка данных проводились командой Sber AI и компанией SberDevices.

Коллектив авторов: Игорь Павлов, Юлия Агафонова, Соня Кириллова, Михаил Шойтов, Владимир Архипкин, Зейн Шахин, Виктория Вульф, Константин Шуршуков, Сергей Марков, Андрей Кузнецов и Денис Димитров.

© Habrahabr.ru