Генерация аниме с помощью нейросети StyleGAN

StyleGAN — это нейросеть, генерирующая лица никогда не существовавших людей. Гверн Бранвен (Gwern Branwen), исследователь нейросетей из США, смог вывести ее на новый уровень — StyleGAN научилась создавать лица несуществующих персонажей аниме. Также он запустил свой сайт, на котором каждые 18 секунд появляется новый персонаж и сюжет аниме, сгенерированный GPT-3. В этой статье мы расскажем, как самостоятельно создать генератор аниме-персонажей, используя StyleGAN.

Немного истории

Генерация аниме-лиц в хорошем качестве — задача, которая долгое время не поддавалась искусственному интеллекту. Несмотря на успехи ИИ в создании человеческих лиц, аниме-изображения выходили расплывчатыми и похожими друг на друга. Гверн пытался обучить целый ряд генеративных нейросетей: StackGAN / StackGAN ++ & Pixel * NN *, WGAN-GP, Glow, GAN-QP, MSG-GAN, SAGAN, VGAN, PokeGAN, BigGAN 3, ProGAN и StyleGAN. Большинство попыток провалились: сети расходились после 1-2 дней обучения, создавали одних и тех же персонажей, или же просто выдавали рисунки плохого качества.

Хорошие результаты продемонстрировали нейросети BigGAN и ProGAN: они смогли создать и масштабировать четкие аниме-лица. Однако обучение ProGAN требовало много затрат и времени (около 6 GPU-недель). Поэтому решено было выбрать StyleGAN — более быструю архитектуру, на которой можно натренировать большую модель на большем наборе данных.

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

Примеры изображений

Рандомный персонаж аниме, созданный нейросетью:

Примеры аниме-персонажей, сгенерированных StyleGAN:

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

О StyleGAN

StyleGAN — генеративно-состязательная сеть, представленная в 2018 году. Она использует архитектуру GAN, которая применяется и в ProGAN, но черпает вдохновение из механизма передачи стиля. StyleGAN модифицирует свою генераторную сеть (генератор), которая создает изображение путем его многократного увеличения: 8px → 16px → 32px → 64px → 128px и так далее. На каждом уровне используется комбинация случайных входных данных — «стилевого шума» («style noise») в сочетании с AdaIN. Он сообщает генератору, как стилизовать изображения в определенном разрешении: изменить волосы, текстуру кожи и так далее. «Стилевой шум» при низком разрешении (таком как 32px) влияет на картинку глобально — к примеру, может определить длину волос, а при высоком разрешении (таком как 256px) влияет на то, как выглядят отдельные пряди. Систематически обеспечивая некоторую случайность на каждом этапе процесса создания изображения, StyleGAN может эффективно выбирать наиболее удачные варианты.

StyleGAN вносит ряд дополнительных улучшений: например, вводит новый набор данных лиц «FFHQ» с изображениями размером 1024 пикселя (больше, чем у ProGAN). Также StyleGAN демонстрирует меньше потерь и с точки зрения архитектуры необычайно интенсивно использует полносвязные слои для обработки случайного ввода — не менее 8 слоев по 512 нейронов, тогда как большинство GAN используют 1 или 2 слоя. Ещё более поразительным является то, что StyleGAN не использует методы, которые другие GAN считают критически важными для возможности обучения в масштабе 512-1024 пикселей: релятивистские потери, распределение шума, расширенная регуляризация и так далее.

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

Применение StyleGAN

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

• «Этого человека не существует»,
• тесты: какое лицо настоящее, реальное или поддельное,
• генерация инстаграм-портретов,
• коты: «Этих кошек не существует», «Этого кота не существует» (ошибки в генерации кошек; интерполяция / передача стиля),
• гостиничные номера (с текстовыми описаниями, сгенерированными char-RNN): «Этого номера не существует»,
• «Этой вайфу не существует»,
• шрифты,
• готические соборы,
• спутниковые снимки,
• граффити,
• городские пейзажи,
• селфи из фотобудки,
• рамен,
• винные этикетки,
• портреты,
• покемоны,
• логотипы,
• текстуры Doom,
• вазы,
• футболки,
• бабочки.

Что необходимо для обучения

Данные

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

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

Вычисление

Чтобы обеспечить достаточный размер мини-пакетов, вам потребуются графические процессоры с >11 ГБ видеопамяти. К примеру, для этого может подойти Nvidia 1080ti или Nvidia Tesla V100.

В репозитории StyleGAN указано расчетное время обучения для систем с 1–8 графическими процессорами. Сконвертировав их в общее количество часов работы графического процессора, мы получим следующее:

Подготовка данных

Самая сложная часть запуска StyleGAN — правильная подготовка набора данных. StyleGAN, в отличие от большинства реализаций GAN (особенно PyTorch), не поддерживает чтение каталога файлов в качестве входных данных. Он может только читать свой уникальный формат .tfrecord, в котором каждое изображение хранится в виде необработанных массивов для каждого соответствующего ему разрешения. Таким образом, исходные данные должны быть полностью однородными, преобразованными в формат .tfrecord специальным скриптом dataset_tool.py. При этом они будут занимать примерно в 19 раз больше места на диске.

Важно: набор данных StyleGAN должен состоять из изображений, отформатированных одинаково. Картинки должны быть только одинакового формата — 512×512 пикселей, 1024×1024 пикселей и так далее. Также они должны относиться к одному цветовому пространству (нельзя использовать одновременно RGB и градации серого). Изображения не должны быть прозрачными. Тип файла должен быть таким же, как у модели, которую вы собираетесь обучить/переобучить (нельзя переобучить PNG-модель на наборе данных JPG). Не должно быть скрытых ошибок, таких как ошибки контрольной суммы CRC, которые часто игнорируют программы просмотра изображений или библиотеки.

Подготовка лиц

Процесс подготовки лиц выглядит следующим образом:

1. Скачайте необработанные изображения с Danbooru2018.
2. Извлеките из метаданных JSON Danbooru2018 все идентификаторы подмножества изображений, если требуется определенный тег Danbooru (например, один символ), используя jq и shell-сценарии.
3. Обрежьте лица персонажей на изображениях, используя метод lbpcascade_animeface от Nagadomi (обычные методы обнаружения лиц не работают на аниме).
4. Удалите пустые, монохромные или чёрно-белые картинки, а также дубликаты.
5. Конвертируйте изображения в JPG.
6. Уменьшить разрешение до необходимого размера (512 пикселей) при помощи waifu2x.
7. По возможности улучшите качество данных, проверяя изображения низкого качества вручную. Удалите и отфильтруйте с помощью предварительно обученной GAN почти полные дубликаты, найденные с помощью findimagedupes.
8. Преобразуйте данные в формат StyleGAN с помощью dataset_tool.py.

Ваша цель — превратить это:

в это:

Ниже представлены shell-скрипты для подготовки набора данных. В качестве альтернативы можно использовать утилиту danbooru-utility, которая помогает исследовать набор данных, фильтровать по тегам, оценивать, обнаруживать лица и изменять размер изображений.

Обрезка

Загрузить набор данных Danbooru2018 можно через BitTorrent или rsync. Вы получите тарбол метаданных JSON, распакованный в папку metadata/2* и структуру данных {original,512px}/{0-999}/$ID.{png,jpg,...}.

Идентификаторы можно извлечь из имен файлов напрямую или из метаданных:

После установки и тестирования Nagadomi lbpcascade_animeface можно использовать простой скрипт для обрезки лиц — crop.py. Точность на изображениях Danbooru довольно хорошая: примерно 90% идеальных лиц, 5 % некачественных лиц и 5 % ошибочных, на которых вместо лиц изображены другие части тела. Точность можно улучшить, если сделать сценарий более строгим. Например, можно увеличить область интереса до размера 250×250 пикселей.

Этот скрипт может обработать лишь одно входное изображение. Обрезать все изображения параллельно можно следующим образом:

Очистка и масштабирование

Удаление ненужных файлов происходит следующим образом:

## Удаление ошибочных и пустых файлов:

find faces/ -size 0 -type f -delete 



## Удаление слишком маленьких файлов: 

find faces/ -size -40k -type f -delete 



## Удаление дубликатов: 

fdupes --delete --omitfirst --noprompt faces/ 



## Удаление монохромных или малоцветных изображений: эвристика <257 

### уникальных цветов несовершенна, но лучше всего, что было испробовано 

deleteBW() { if [[ `identify -format "%k" "$@"` -lt 257 ]]; 

 then rm "$@"; fi; } 

export -f deleteBW 

find faces -type f | parallel --progress deleteBW 

Следующим важным шагом является масштабирование картинок с помощью waifu2x⁠, которая отлично справляется с двукратным увеличением аниме-изображений. Недостатком является то, что она работает довольно медленно (1-10 секунд на изображение), запускается только на GPU и написана в неподдерживаемом на данный момент DL-фреймворке Torch.

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

Использование waifu2x:

. ~/src/torch/install/bin/torch-activate
upscaleWaifu2x() { 

    SIZE1=$(identify -format "%h" "$@") 

    SIZE2=$(identify -format "%w" "$@"); 



    if (( $SIZE1 < 512 && $SIZE2 < 512 )); then 

        echo "$@" $SIZE 

        TMP=$(mktemp "/tmp/XXXXXX.png") 

        CUDA_VISIBLE_DEVICES="$((RANDOM % 2 < 1))" nice th \ 

             ~/src/waifu2x/waifu2x.lua -model_dir \ 

             ~/src/waifu2x/models/upconv_7/art -tta 1 -m noise_scale \ 

            -noise_level 1 -scale 2 \ 

            -i "$@" -o "$TMP" 

        convert "$TMP" "$@" 

        rm "$TMP" 

    fi; } 



export -f upscaleWaifu2x

find faces/ -type f | parallel --progress --jobs 9 upscaleWaifu2x

Проверка и дополнение данных

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

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

Повторное масштабирование и преобразование

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

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

Далее масштабируем все изображения до размера 512х512 пикселей с помощью ImageMagick без сохранения пропорций:

Любое немного отличающееся изображение может привести к сбою процесса импорта. Поэтому удаляем все изображения, которые хоть немного отличаются от 512×512 sRGB JPG:

Теперь можно преобразовать данные в формат StyleGAN с помощью dataset_tool.py. Настоятельно рекомендуем отредактировать dataset_tool.py и добавить в этот скрипт вывод имен файлов. Если произойдет сбой, то вы сможете определить, из-за какого изображения он произошел:

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

Обучение модели

Установка

Предполагается, что у вас установлена и работает CUDA. Автор использовал ОС Ubuntu Bionic 18.04.2 LTS, драйвер Nvidia версии № 410.104, CUDA 10.1 и TensorFlow 1.13.1.

Помимо этого, вам понадобится Python версии 3.6 и выше, TensorFlow и зависимости для StyleGAN:

Конфигурация

StyleGAN не поддерживает опции командной строки. Вместо этого нужно отредактировать train.py и train/training_loop.py:

Исходная конфигурация указана в функции training_loop в строке 112.

Ключевые аргументы:

• G_smoothing_kimg и D_repeats — влияют на динамику обучения,
• network_snapshot_ticks — частота сохранения скриншотов,
• resume_run_id — установлен на "latest",
• resume_kimg — указывает на финальное значение разрешения создаваемых изображений.

Рекомендуется установить minibatch_repeats = 1 вместо minibatch_repeats = 5. Это делает обучение более стабильным и быстрым.

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

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

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

Мы установили разрешение изображений размером 512 пикселей для набора данных dataset/faces, включили зеркальное отражение и задали заголовок для контрольных точек и логов, который теперь появится в results/ в строке '-faces'.

Если у вас нет 8 графических процессоров, вам необходимо изменить значение -preset и указать количество ваших GPU. StyleGAN не станет автоматически выбирать правильное количество графических процессоров. Он попытается использовать несуществующие графические процессоры и после выдаст сообщение об ошибке. Обратите внимание, что в CUDA используется нулевое индексирование, то есть GPU:0 указывает на первый процессор, GPU:1 на второй и так далее. Как выглядит сообщение об ошибке:

Для 2×1080ti автор установил:

Запуск

StyleGAN можно запустить в GNU Screen и разделить его на несколько оболочек: 1 терминал для запуска StyleGAN, 1 для TensorBoard и 1 для Emacs.

В Emacs можно открыть файлы train.py и train/training_loop.py для справки и удобного редактирования.

Последний патч для StyleGAN позволяет включить цикл while, чтобы нейросеть продолжала работу после сбоев, например:

TensorBoard — это утилита, позволяющая визуализировать записанные в ходе обучения значения, которые можно просматривать в браузере, например:

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

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

Вот так выглядит процесс успешного обучения:
https://gwern.net/image/gan/stylegan/2019-03-16-stylegan-facestraining.mp4

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

После успешного обучения StyleGAN приступаем к самому интересному — созданию образцов!

Трюк с усечением

Трюк с усечением Ψ является наиболее важным гиперпараметром для всего поколения StyleGAN. Он используется во время генерации выборки, но не во время обучения. С помощью трюка усечения можно контролировать разнообразие генерируемых изображений. При Ψ = 0 разнообразие равно нулю, и все лица представляют собой одно глобальное усредненное лицо. При ±0,5 генерируется широкий диапазон лиц, а при ±1,2 вы увидите огромное разнообразие лиц и стилей, но также получите много искажений. Рекомендуется установить Ψ = ±0,7, это обеспечивает наилучший баланс качества и разнообразия картинок.

Случайные образцы

В репозитории StyleGAN есть скрипт pretrained_example.py для загрузки и создания одного лица. Однако его можно легко адаптировать для использования локальной модели и генерировать, допустим, 1000 выборочных изображений с гиперпараметром Ψ = 0,6, что даст качественные, но не очень разнообразные изображения. Они будут сохраняться в results/example-{0-999}.png:

Видео

Самый простой способ получить видео процесса обучения — сделать скриншоты прогресса и конвертировать их в видеозапись. В ходе обучения размер скриншотов растет по мере увеличения качества и детализации, и в конце может быть довольно большим. Поэтому стоит выполнить сжатие с потерями или преобразование в JPG. Чтобы превратить множество снимков в обучающее видео, можно использовать FFmpeg:

Загрузка готовой модели

Вы можете загрузить готовую модель StyleGAN с размером изображений 512 пикселей, обученную на 218794 лицах в течение примерно 38 GPU-дней.

Перенос обучения

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

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

Основная проблема заключается в том, что повторный запуск обучения не может начинаться с нулевой итерации. В этом случае StyleGAN игнорирует предварительно обученную модель и запускает обучение заново. Чтобы этого избежать, сразу укажите высокое разрешение. Например, чтобы начать с размера изображения в 512 пикселей, установите resume_kimg=7000 в training_loop.py. Тогда StyleGAN будет вынужден загрузить предыдущую модель. Чтобы убедиться, что вы всё сделали правильно, до переноса обучения проверьте первый образец (fakes07000.png или какой-нибудь ещё): он должен выглядеть как исходная модель в конце обучения. Затем последующие обучающие образцы должны показать, что оригинал быстро трансформируется в новый набор данных.

Так выглядит перенос обучения StyleGAN на конкретного персонажа — Холо из аниме «Волчица и пряности»:

Другой пример — Сорью Аска Лэнгли из «Евангелиона»:
https://gwern.net/image/gan/stylegan/2019-02-21-stylegan-asukatransfer-interpolation.mp4

Теперь, если вам понадобится портрет в стиле аниме или уникальная аватарка, то вы сможете сгенерировать ее самостоятельно. А на этом сайте вы можете сделать из фото аниме — тоже с помощью нейросети!

С оригинальной статьей можно ознакомиться по ссылке.