Погружаемся в StyleGAN

Или как я лишил StSpyro работы

Одним осенним днём меня посреди ночи осенило, что я хочу научиться делать очень красивые картинки с помощью StyleGAN. Подумал, что меня останавливает заняться этим прямо сейчас. И занялся. В этой статье я опишу все перепитии, с которыми я столкнулся по пути, а ещё покажу конечный результат.

Прежде всего, расскажу вам о StyleGAN, больше всего ориентируясь на этот трактат (он на английском, но если вы его понимаете, всем крайне советую почитать). Как несложно предположить, название StyleGAN состоит из двух компонент: Style и GAN.

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

Style исходит из техники "style transfer", которую раньше применяли для переноса стиля с одних изображений на другие. Разработчикам StyleGAN показалось, что это будет хорошей идеей - попытаться той же техникой переносить стиль на случайный шум и тем самым генерировать изображения с нуля, и они, в общем-то, оказались правы настолько, что мир до сих пор расхлёбывает последствия.

Генератор в StyleGAN работает по свёрточной модели: в нейронной сети генератора находятся несколько слоёв, каждый из которых больше предыдущего (то есть, первый 4х4, второй 8х8, третий 16х16 и так далее до желаемого разрешения картинки) и каждый из которых последовательно за предыдущим уточняет изображение, подкрашивая его и добавляя в него новые детали. К примеру, первый слой задаёт фон и цвет центрального объекта, второй уточняет цвета и добавляет основу для деталей, третий дорисовывает детали, а четвёртый добавляет чёткости и узоров.

StyleGAN в отличии от обычной реализации GAN для изображений добавляет "артистичный" шум на каждом вычислительном слое, а не только на входе самого первого слоя. Таким образом, у него лучше и чётче получаются мелкие детали, которые должны дорисовываться на последних слоях генератора.

На данный момент существует уже четыре реализации StyleGAN: StyleGAN, StyleGAN2, StyleGAN2-ADA и StyleGAN3. Каждая из них является логичным продолжением предыдущей, решая определённые проблемы в генерациях, но в целом они все работают по одной вышеописанной системе.

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

Я в свою очередь пользуюсь бесплатным Google Colab (бесплатный, потому что платный аккаунт нельзя купить в СНГ), который даёт мне раз в сутки на 4 часа NVIDIA Tesla K80. Полноценное обучение с таким оборудованием займёт у вас не более 100 часов.

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

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

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

Я тренировал нейросеть на картинках с Нахтигалкой. Основной сет брал из своих колажей, колажей StSpyro, а ещё нагло спёр вот эти стикерпаки. Все их я при необходимости обрезал и свёл к разрешению 256х256. Всего у меня в базе картинок оказалось 165 изображений.

Сначала я тренировал модель с нуля, но потом появилась необходимость перезапустить вычисления (я перешёл на другой фреймворк и старая модель оказалась несовместима) и я начал переобучать модель ffhq1024 (та самая с thispersondoesnotexist.com), отчего прогресс пошёл в несколько раз быстрее. Видимо, коты не сильно отличаются от людей.

Мне понадобилось около двух сессий по 4-6 часов для осмысленных результатов, далее я медленно дообучал сеть, поскольку чем дальше обучение, тем меньше видим прогресс: изображения просто становятся более чётче, а процент удачных результатов растёт.

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

Коллаж из результатов вы можете увидеть ниже:

Особенно удачные случаи из колажа выше, по группам

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

Всем спасибо за внимание!