Почему StyleGAN превосходит другие модели для генерации лиц и как тренировать свой собственный StyleGAN

Генеративные состязательные сети (GAN) - это архитектура, представленная Яном Гудфеллоу и его коллегами в 2014 году для генеративного моделирования, которая использует модель для создания новых выборок, имитирующих существующий набор данных. Он состоит из двух сетей: генератора, который генерирует новые образцы, и дискриминатора, который обнаруживает поддельные образцы. Генератор пытается обмануть дискриминатор, в то время как дискриминатор пытается обнаружить образцы, синтезированные генератором.

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

Сегодня GAN бывают разных форм: DCGAN, CycleGAN, SAGAN… Из множества GAN на выбор я использовал StyleGAN для создания искусственных лиц знаменитостей. В этом отчете я объясню, что делает архитектуру StyleGAN хорошим выбором, как обучать модель и некоторые результаты обучения.

Архитектура модели StyleGAN

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

В 2018 году NVIDIA предложила новую сеть в своей статье Архитектура на основе стилей для GAN. Предыдущие модели GAN уже показали способность генерировать человеческие лица, но одна проблема заключается в том, чтобы иметь возможность управлять некоторыми особенностями сгенерированных изображений, такими как цвет волос или поза. StyleGAN пытается решить эту проблему путем включения и развития прогрессивного обучения для изменения каждого уровня детализации отдельно. При этом он может управлять визуальными особенностями, выраженными на каждом уровне детализации, от грубых особенностей, таких как поза и форма лица, до более мелких деталей, таких как цвет глаз и форма носа, не затрагивая другие уровни.

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

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

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

Обработка данных

Я использовал набор данных CelebA для обучения моей модели. CelebA содержит 202 599 изображений лиц 10 177 разных знаменитостей. Исходный набор данных аннотирован двоичными элементами, такими как очки или большой нос, но мы будем использовать только сами изображения для создания лиц.

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

Мы изменяем размер изображений, чтобы приспособиться к прогрессивному обучению, как описано ранее, с 8x8 до 128x128, выбранного окончательного размера вывода. Обратите внимание, что этот метод позволяет обучать модели с изображениями 1024x1024, но для этого потребуется более месяца обучения графическому процессору, а разрешение не менее 64x64 уже дает хорошие визуальные результаты. Каждое изображение изменяется, чтобы иметь копию в размерах 8x8, 16x16, 32x32, 64x64 и 128x128, так что обученный генератор будет генерировать изображения размером 128x128.

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

Обучение

Для обучения собственной модели я нашел на Github отличную реализацию StyleGAN в моем любимом фреймворке машинного обучения с понятным кодом.

Обучение проводится так же, как и в традиционных сетях GAN, с добавленной задачей прогрессивного обучения. Я использовал оптимизатор Adam со скоростью обучения 0,001 как для дискриминатора, так и для генератора. Я использую размер пакета 16 из-за ограничений памяти и размер кода 512, то есть вектор случайного шума, вводимый в генератор, имеет размер 1x512. Для функции потерь я использую потери Вассерштейна. Для прогрессивного обучения каждый размерный размер обучается на 600 000 экземпляров изображений перед увеличением размера, пока не достигнет размера 128x128, где я продолжаю тренировать модель до сходимости.

Мы можем использовать три метода регуляризации. Во-первых, при загрузке входных изображений во время обучения выполняется переворот по горизонтали с вероятностью 0,5. Таким образом, мне не нужно увеличивать размер набора данных, что замедлит обучение, но модель встречает больше различных входных данных во время обучения.

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

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

Генерация изображений

Мы используем только обученный генератор для генерации изображений. После обучения модели мы сначала вычисляем средний вектор признаков. Это делается просто путем использования генератора для синтеза нескольких изображений, в нашем случае мы использовали 10, а затем для вычисления среднего значения векторов признаков, как упоминалось ранее, по полученным изображениям для каждого уровня. Затем во время генерации используется вектор признаков v_mean + ψ (v - v_mean), где v - это выходные данные сети сопоставления признаков, v_mean - это ранее вычисленный средний вектор признаков, а ψ - константа, которая контролирует силу среднего вектора признаков. Чем больше c, тем больше окончательное сгенерированное изображение будет отличаться от «среднего» лица. Мы даже можем использовать разные значения ψ для каждого уровня, чтобы контролировать, насколько далеко от среднего значения каждой функции.

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

Полученные результаты

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

Итерация 100, размерность 8x8

Итерация 10000, размер 8х8

Итерация 50000, размерность 16x16

Итерация 100000, размерность 32x32

Итерация 140000, размерность 64x64

Итерация 155000, размерность 128х128

Итерация 400000, размерность 128х128

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

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

Спасибо за чтение, и я надеюсь, что вам понравилась статья! Не стесняйтесь сообщать мне о любых интересных экспериментах, которые вы проводите со StyleGAN :)

Вы можете найти полный код, который я использовал, от загрузки набора данных CelebA до обучения и генерации примеров с использованием StyleGAN на моем Github, но полная заслуга за реализацию StyleGAN в PyTorch идет на rosinality.