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

StyleGAN — та самая нейросеть, которая генерирует лица несуществующих людей на сайте thispersondoesnotexist.com. Исследователь Gwern Branwen вывел её на новый уровень и научил создавать несуществующие лица персонажей аниме, запустив свой сайт thiswaifudoesnotexist.net. На нём каждые 15 секунд появляется новый персонаж и его история, также генерируемая искусственным интеллектом. Сегодня мы расскажем, как самому обучить нейросеть для создания аниме-лиц с помощью StyleGAN.

В попытках заставить компьютер рисовать аниме есть что-то забавное — это уж точно интереснее, чем работать со снимками знаменитостей или с датасетом ImageNet! Кроме того, рисунки и аниме отличаются от фотографий, которые сейчас очень часто используются в задачах машинного обучения. Если нейросети научатся генерировать случайные нарисованные изображения, то совсем скоро станет возможна генерация картин и фотографий по их текстовому описанию.

Gwern пробовал решить задачу с помощью ряда существующих генеративных сетей: StackGAN / StackGAN ++ & Pixel * NN *, WGAN-GP, Glow, GAN-QP, MSG-GAN, SAGAN, VGAN, PokeGAN, BigGAN 3, ProGAN и StyleGAN. Многие из них начинали либо расходиться после одного или двух дней обучения, либо создавать ограниченный диапазон лиц (или одно лицо), либо просто сходиться к изображениям плохого качества.

Первый успех продемонстрировали сети BigGAN и ProGAN: они показали, что обычные CNN могут научиться генерировать и масштабировать чёткие аниме-изображения. Сеть ProGAN была достаточно мощной, но требовала около 6 недель обучения на GPU. Поэтому ей на смену пришла StyleGAN — более быстрая архитектура, с которой можно обучать объёмные модели на больших наборах данных.

Аниме — *Сгенерированные с помощью GAN аниме*

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

Примеры

Для начала продемонстрируем возможность StyleGAN генерировать аниме-лица:

озможность StyleGAN генерировать аниме-лица:

64 Ð¸Ð· Ð»ÑÑÑÐ¸Ñ Ð¾Ð±ÑÐ°Ð·ÑÐ¾Ð² Ð°Ð½Ð¸Ð¼Ðµ TWDNE, Ð²ÑÐ±ÑÐ°Ð½Ð½ÑÑ Ð¸Ð· ÑÐ¾ÑÐ¸Ð°Ð»ÑÐ½ÑÑ ÑÐµÑÐµÐ¹ (Ð½Ð°Ð¶Ð¼Ð¸ÑÐµ, ÑÑÐ¾Ð±Ñ ÑÐ²ÐµÐ»Ð¸ÑÐ¸ÑÑ). — *64 лучших образца аниме-лиц TWDNE (This Waifu Does Not Exist), собранных из социальных сетей*

100 random sample images from the StyleGAN anime faces on TWDNE — *100 случайно сгенерированных StyleGAN аниме-изображений*

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

Видео, демонстрирующие разнообразие аниме-лиц:

StyleGAN anime face interpolation videos are Elon Muskâ¢-approved! — *Илон Маск одобряет*

О StyleGAN

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

Karras et al 2018, StyleGAN vs ProGAN architecture: âFigure 1. While a traditional generator [29] feeds the latent code [z] though the input layer only, we first map the input to an intermediate latent space W, which then controls the generator through adaptive instance normalization (AdaIN) at each convolution layer. Gaussian noise is added after each convolution, before evaluating the nonlinearity. Here âAâ stands for a learned affine transform, and âBâ applies learned per-channel scaling factors to the noise input. The mapping network f consists of 8 layers and the synthesis network g consists of 18 layersâtwo for each resolution (42-â10242). The output of the last layer is converted to RGB using a separate 1Ã1 convolution, similar to Karras et al. [29]. Our generator has a total of 26.2M trainable parameters, compared to 23.1M in the traditional generator." — *Сравнение архитектур ProGAN (a) и StyleGAN (b)*

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

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

Приложения

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

— Этого человека не существует

тесты: Какое лицо реально / Реальное или нет?

— котики: Этих кошек не существует , Этого кота не существует (ошибки в генерации кошек ; интерполяция / передача стиля)

— гостиничные номера (с текстовыми описаниями, сгенерированными RNN): Этого номера не существует

кухня / столовая / гостиная / спальня (с использованием трансферного обучения)

— Этого вайфу не существует

— шрифты

— рамен

— художественные портреты

— покемоны

— логотипы

— текстуры Doom

Imagequilt visualization of the wide range of visual subjects StyleGAN has been applied to — *Широкий спектр возможностей StyleGAN*

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

Данные

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

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

Вычисление

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

Репозиторий StyleGAN предоставляет приблизительные сроки обучения для систем с 1-8 GPU. Сконвертировав их в общее количество часов работы GPU, мы получим следующее:

Предполагаемое время обучения StyleGAN при различных разрешениях изображений и используемых GPU (источник: репозиторий StyleGAN)

Число используемых GPU	1024 ²	512 ²	256 ²	[Апрель 2019 г., тарифы REG.RU ]
1	41 день 4 часа [988 GPU-часов]	24 дня 21 час [597 GPU-часов]	14 дней 22 часа [358 GPU-часов]	[₽89к, ₽54к, ₽32к]
2	21 день 22 часа [1052]	13 дней 7 часов [638]	9 дней 5 часов [442]	[105к, 64к, 44к]
4	11 дней 8 часов [1088]	7 дней 0 часов [672]	4 дня 21 час [468]	[131к, 81к, 56к]
8	6 дней 14 часов [1264]	4 дня 10 часов [848]	3 дня 8 часов [640]	[177к,, 119к, 90к]

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

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

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

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

Процесс подготовки лиц:

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

Цель состоит в том, чтобы превратить это:

в это:

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

Обрезка

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

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

find ./512px/ -type f | sed -e 's/.*\/\([[:digit:]]*\)\.jpg/\1/'
# 967769
# 1853769
# 2729769
# 704769
# 1799769
# ...
tar xf metadata.json.tar.xz
cat metadata/* | jq '[.id, .rating]' -c | fgrep '"s"' | cut -d '"' -f 2 # "
# ...

find ./512px/ -type f | sed -e 's/.*\/$[[:digit:]]*$\.jpg/\1/'

# 967769

# 1853769

# 2729769

# 704769

# 1799769

# ...

tar xf metadata.json.tar.xz

cat metadata/* | jq '[.id, .rating]' -c | fgrep '"s"' | cut -d '"' -f 2 # "

# ...

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

crop.py:

import cv2
import sys
import os.path

def detect(cascade_file, filename, outputname):
    if not os.path.isfile(cascade_file):
        raise RuntimeError("%s: not found" % cascade_file)

    cascade = cv2.CascadeClassifier(cascade_file)
    image = cv2.imread(filename)
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    gray = cv2.equalizeHist(gray)

    ## Предлагаемая модификация: увеличьте minSize до '(250,250)' пикс.
    ## Это повысит число качественных лиц и уменьшит количество ложных
    ## срабатываний. Лица размером всего лишь 50x50px обычно либо 
    ## бесполезны, либо вообще не лица.

    faces = cascade.detectMultiScale(gray,
    # Параметры детектора
    scaleFactor = 1.1,
    minNeighbors = 5,
    minSize = (50, 50))
    i=0
    for (x, y, w, h) in faces:
        cropped = image[y: y + h, x: x + w]
        cv2.imwrite(outputname+str(i)+".png", cropped)
        i=i+1

if len(sys.argv) != 4:
    sys.stderr.write("usage: detect.py <animeface.xml file> <input> \
    <output prefix>\n")
    sys.exit(-1)

detect(sys.argv[1], sys.argv[2], sys.argv[3])

import cv2

import sys

import os.path

def detect(cascade_file, filename, outputname):

if not os.path.isfile(cascade_file):

raise RuntimeError("%s: not found" % cascade_file)

cascade = cv2.CascadeClassifier(cascade_file)

image = cv2.imread(filename)

gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

gray = cv2.equalizeHist(gray)

## Предлагаемая модификация: увеличьте minSize до '(250,250)' пикс.

## Это повысит число качественных лиц и уменьшит количество ложных

## срабатываний. Лица размером всего лишь 50x50px обычно либо

## бесполезны, либо вообще не лица.

faces = cascade.detectMultiScale(gray,

# Параметры детектора

scaleFactor = 1.1,

minNeighbors = 5,

minSize = (50, 50))

i=0

for (x, y, w, h) in faces:

cropped = image[y: y + h, x: x + w]

cv2.imwrite(outputname+str(i)+".png", cropped)

i=i+1

if len(sys.argv) != 4:

sys.stderr.write("usage: detect.py <animeface.xml file> <input> \

<output prefix>\n")

sys.exit(-1)

detect(sys.argv[1], sys.argv[2], sys.argv[3])

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

cropFaces() { 
    BUCKET=$(printf "%04d" $(( $@ % 1000 )) )
    ID="$@"
    CUDA_VISIBLE_DEVICES="" nice python \
    ~/src/lbpcascade_animeface/examples/crop.py \
    ~/src/lbpcascade_animeface/lbpcascade_animeface.xml \
    ./original/$BUCKET/$ID.* "./faces/$ID"
}
export -f cropFaces

mkdir ./faces/
cat sfw-ids.txt | parallel --progress cropFaces

cropFaces() {

BUCKET=$(printf "%04d" $(( $@ % 1000 )) )

ID="$@"

CUDA_VISIBLE_DEVICES="" nice python \

~/src/lbpcascade_animeface/examples/crop.py \

~/src/lbpcascade_animeface/lbpcascade_animeface.xml \

./original/$BUCKET/$ID.* "./faces/$ID"

}

export -f cropFaces

mkdir ./faces/

cat sfw-ids.txt | parallel --progress cropFaces

Здесь мы устанавливаем параметр CUDA_VISIBLE_DEVICES="" — это отключает использование GPU, поскольку для этой задачи он не даёт очевидного ускорения.

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

Удаление ненужных файлов:

<pre class="wp-block-preformatted nums:false">## Удаление ошибочных и пустых файлов:
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 </pre>

<pre class="wp-block-preformatted nums:false">## Удаление ошибочных и пустых файлов:

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 </pre>

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

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

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

<pre class="wp-block-preformatted nums:false">. ~/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</pre>

<pre class="wp-block-preformatted nums:false">. ~/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</pre>

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

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

dataAugment () {
    image="$@"
    target=$(basename "$@")
    suffix="png"
    convert -deskew 50 "$image" "$target".deskew."$suffix"
    convert -resize 110%x100% "$image" "$target".horizstretch."$suffix"
    convert -resize 100%x110% "$image" "$target".vertstretch."$suffix"
    convert -blue-shift 1.1 "$image" "$target".midnight."$suffix"
    convert -fill red -colorize 5% "$image" "$target".red."$suffix"
    convert -fill orange -colorize 5% "$image" "$target".orange."$suffix"
    convert -fill yellow -colorize 5% "$image" "$target".yellow."$suffix"
    convert -fill green -colorize 5% "$image" "$target".green."$suffix"
    convert -fill blue -colorize 5% "$image" "$target".blue."$suffix"
    convert -fill purple -colorize 5% "$image" "$target".purple."$suffix"
    convert -adaptive-blur 3x2 "$image" "$target".blur."$suffix"
    convert -adaptive-sharpen 4x2 "$image" "$target".sharpen."$suffix"
    convert -brightness-contrast 10 "$image" "$target".brighter."$suffix"
    convert -brightness-contrast 10x10 "$image" \
    "$target".brightercontraster."$suffix"
    convert -brightness-contrast -10 "$image" "$target".darker."$suffix"
    convert -brightness-contrast -10x10 "$image" \
    "$target".darkerlesscontrast."$suffix"
    convert +level 5% "$image" "$target".contraster."$suffix"
    convert -level 5%\! "$image" "$target".lesscontrast."$suffix"
 }
export -f dataAugment
find faces/ -type f | parallel --progress dataAugment

dataAugment () {

image="$@"

target=$(basename "$@")

suffix="png"

convert -deskew 50 "$image" "$target".deskew."$suffix"

convert -resize 110%x100% "$image" "$target".horizstretch."$suffix"

convert -resize 100%x110% "$image" "$target".vertstretch."$suffix"

convert -blue-shift 1.1 "$image" "$target".midnight."$suffix"

convert -fill red -colorize 5% "$image" "$target".red."$suffix"

convert -fill orange -colorize 5% "$image" "$target".orange."$suffix"

convert -fill yellow -colorize 5% "$image" "$target".yellow."$suffix"

convert -fill green -colorize 5% "$image" "$target".green."$suffix"

convert -fill blue -colorize 5% "$image" "$target".blue."$suffix"

convert -fill purple -colorize 5% "$image" "$target".purple."$suffix"

convert -adaptive-blur 3x2 "$image" "$target".blur."$suffix"

convert -adaptive-sharpen 4x2 "$image" "$target".sharpen."$suffix"

convert -brightness-contrast 10 "$image" "$target".brighter."$suffix"

convert -brightness-contrast 10x10 "$image" \

"$target".brightercontraster."$suffix"

convert -brightness-contrast -10 "$image" "$target".darker."$suffix"

convert -brightness-contrast -10x10 "$image" \

"$target".darkerlesscontrast."$suffix"

convert +level 5% "$image" "$target".contraster."$suffix"

convert -level 5%\! "$image" "$target".lesscontrast."$suffix"

}

export -f dataAugment

find faces/ -type f | parallel --progress dataAugment

Масштабирование (ещё раз) и преобразование

После дополнения данных неплохо было бы сэкономить место на диске, преобразовав изображения в JPG и снизив качество примерно на 33%. Так можно освободить достаточно много дискового пространства без видимых потерь:

convertPNGToJPG() { convert -quality 33 "$@" "$@".jpg &amp;&amp; rm "$@"; }
export -f convertPNGToJPG
find faces/ -type f -name "*.png" | parallel --progress convertPNGToJPG

convertPNGToJPG() { convert -quality 33 "$@" "$@".jpg && rm "$@"; }

export -f convertPNGToJPG

find faces/ -type f -name "*.png" | parallel --progress convertPNGToJPG

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

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

find faces/ -type f | xargs --max-procs=16 -n 5000 \
mogrify -resize 512x512\&gt; -extent 512x512\&gt; -gravity center -background black

1 2	find faces/ -type f \| xargs --max-procs=16 -n 5000 \ mogrify -resize 512x512\> -extent 512x512\> -gravity center -background black

Любое ошибочное изображение может привести к сбою процесса импорта, поэтому мы удалим всё, что отличается от JPG 512x512 sRGB:

find faces/ -type f | xargs --max-procs=16 -n 10000 identify | \
    fgrep -v " JPEG 512x512 512x512+0+0 8-bit sRGB"| cut -d ' ' -f 1 | \
    xargs --max-procs=16 -n 10000 rm

find faces/ -type f | xargs --max-procs=16 -n 10000 identify | \

fgrep -v " JPEG 512x512 512x512+0+0 8-bit sRGB"| cut -d ' ' -f 1 | \

xargs --max-procs=16 -n 10000 rm

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

diff --git a/dataset_tool.py b/dataset_tool.py
index 4ddfe44..e64e40b 100755
--- a/dataset_tool.py
+++ b/dataset_tool.py
@@ -519,6 +519,7 @@ def create_from_images(tfrecord_dir, image_dir, shuffle):
    with TFRecordExporter(tfrecord_dir, len(image_filenames)) as tfr:
        order = tfr.choose_shuffled_order()
        if shuffle else np.arange(len(image_filenames))
        for idx in range(order.size):
+           print(image_filenames[order[idx]])
            img = np.asarray(PIL.Image.open(image_filenames[order[idx]]))
            if channels == 1:
            img = img[np.newaxis, :, :] # HW =&gt; CHW

diff --git a/dataset_tool.py b/dataset_tool.py

index 4ddfe44..e64e40b 100755

--- a/dataset_tool.py

+++ b/dataset_tool.py

@@ -519,6 +519,7 @@ def create_from_images(tfrecord_dir, image_dir, shuffle):

with TFRecordExporter(tfrecord_dir, len(image_filenames)) as tfr:

order = tfr.choose_shuffled_order()

if shuffle else np.arange(len(image_filenames))

for idx in range(order.size):

+ print(image_filenames[order[idx]])

img = np.asarray(PIL.Image.open(image_filenames[order[idx]]))

if channels == 1:

img = img[np.newaxis, :, :] # HW => CHW

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

source activate MY_TENSORFLOW_ENVIRONMENT
python dataset_tool.py create_from_images \
    datasets/faces /media/gwern/Data/danbooru2018/faces/

source activate MY_TENSORFLOW_ENVIRONMENT

python dataset_tool.py create_from_images \

datasets/faces /media/gwern/Data/danbooru2018/faces/

Поздравляем, самое сложное позади!

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

Установка

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

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

conda create -n stylegan pip python=3.6
source activate stylegan

## TF:
pip install tensorflow-gpu
## TB:
python -c "import tensorflow as tf; tf.enable_eager_execution(); \
    print(tf.reduce_sum(tf.random_normal([1000, 1000])))"
pip install tensorboard

## StyleGAN:
## Установка зависимостей:
pip install pillow numpy moviepy scipy opencv-python lmdb # requests?
## Загрузка:
git clone 'https://github.com/NVlabs/stylegan.git' &amp;&amp; cd ./stylegan/
## Тестовая установка:
python pretrained_example.py
## Файл ./results/example.png должен быть фотографией мужчины средних лет

conda create -n stylegan pip python=3.6

source activate stylegan

## TF:

pip install tensorflow-gpu

## TB:

python -c "import tensorflow as tf; tf.enable_eager_execution(); \

print(tf.reduce_sum(tf.random_normal([1000, 1000])))"

pip install tensorboard

## StyleGAN:

## Установка зависимостей:

pip install pillow numpy moviepy scipy opencv-python lmdb # requests?

## Загрузка:

git clone 'https://github.com/NVlabs/stylegan.git' && cd ./stylegan/

## Тестовая установка:

python pretrained_example.py

## Файл ./results/example.png должен быть фотографией мужчины средних лет

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

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

1. train/training_loop.py

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

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

G_smoothing_kimg и D_repeats: влияют на динамику обучения

network_snapshot_ticks: указывает, как часто сохранять скриншоты

resume_run_id: установлен на "latest"

resume_kimg: указывает на финальное значение разрешения генерируемых изображений

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

2. train.py (ранее config.py)

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

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

desc += '-faces'; dataset = EasyDict(tfrecord_dir='faces', resolution=512); \
     train.mirror_augment = True

1 2	desc += '-faces'; dataset = EasyDict(tfrecord_dir='faces', resolution=512); \ train.mirror_augment = True

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

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

Так выглядит сообщение об ошибке:

tensorflow.python.framework.errors_impl.InvalidArgumentError: \
Cannot assign a device for operation \
    G_synthesis_3/lod: {{node G_synthesis_3/lod}}was explicitly assigned to \
    /device:GPU:2 but available \
    devices are [ /job:localhost/replica:0/task:0/device:CPU:0, \
    /job:localhost/replica:0/task:0/device:GPU:0, \
    /job:localhost/replica:0/task:0/device:GPU:1, \
    /job:localhost/replica:0/task:0/device:XLA_CPU:0, \
    /job:localhost/replica:0/task:0/device:XLA_GPU:0, \
    /job:localhost/replica:0/task:0/device:XLA_GPU:1 ]. \
    Make sure the device specification refers to a valid device.
      [[{{node G_synthesis_3/lod}}]]

tensorflow.python.framework.errors_impl.InvalidArgumentError: \

Cannot assign a device for operation \

G_synthesis_3/lod: {{node G_synthesis_3/lod}}was explicitly assigned to \

/device:GPU:2 but available \

devices are [ /job:localhost/replica:0/task:0/device:CPU:0, \

/job:localhost/replica:0/task:0/device:GPU:0, \

/job:localhost/replica:0/task:0/device:GPU:1, \

/job:localhost/replica:0/task:0/device:XLA_CPU:0, \

/job:localhost/replica:0/task:0/device:XLA_GPU:0, \

/job:localhost/replica:0/task:0/device:XLA_GPU:1 ]. \

Make sure the device specification refers to a valid device.

[[{{node G_synthesis_3/lod}}]]

У автора для 2x1080ti установлено:

desc += '-preset-v2-2gpus'; submit_config.num_gpus = 2; \
    sched.minibatch_base = 8; sched.minibatch_dict = \
    {4: 256, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8}; \
    sched.G_lrate_dict = {512: 0.0015, 1024: 0.002}; \
    sched.D_lrate_dict = EasyDict(sched.G_lrate_dict); \
    train.total_kimg = 99000

desc += '-preset-v2-2gpus'; submit_config.num_gpus = 2; \

sched.minibatch_base = 8; sched.minibatch_dict = \

{4: 256, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8}; \

sched.G_lrate_dict = {512: 0.0015, 1024: 0.002}; \

sched.D_lrate_dict = EasyDict(sched.G_lrate_dict); \

train.total_kimg = 99000

Запуск

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

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

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

while true; do nice py train.py ; date; (xmessage "alert: StyleGAN crashed" &amp;); \
    sleep 10s; done

1 2	while true; do nice py train.py ; date; (xmessage "alert: StyleGAN crashed" &); \ sleep 10s; done

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

tensorboard --logdir results/02022-sgan-faces-2gpu/
# TensorBoard 1.13.0 at http://127.0.0.1:6006 (Press CTRL+C to quit)

1 2	tensorboard --logdir results/02022-sgan-faces-2gpu/ # TensorBoard 1.13.0 at http://127.0.0.1:6006 (Press CTRL+C to quit)

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

Обучение StyleGAN — это скорее искусство, чем наука. Более подробные советы по настройке и регулировке процесса можно найти на сайте автора проекта.

Так выглядит успешное обучение:

Генерация аниме

А теперь самое интересное — создание образцов!

Для начала поясним значение гиперпараметра 𝜓. Это так называемый «трюк усечения» (“truncation trick”), который используется во время генерации, но не во время обучения. С его помощью можно контролировать разнообразие получаемых изображений. При 𝜓 = 0 разнообразие равно нулю и все лица будут похожи на одно «усреднённое» лицо. При ±0.5 вы получите достаточно широкий диапазон образцов, а при ±1,2 увидите большое разнообразие лиц / стилей, но среди них будет много искажений. Рекомендуется установить значение около ±0.7 — так вы достигните баланса между качеством и разнообразием.

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

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

import os
import pickle
import numpy as np
import PIL.Image
import dnnlib
import dnnlib.tflib as tflib
import config

def main():
    tflib.init_tf()
    _G, _D, Gs = pickle.load( \
    open("results/02051-sgan-faces-2gpu/network-snapshot-021980.pkl", "rb"))
    Gs.print_layers()

    for i in range(0,1000):
        rnd = np.random.RandomState(None)
        latents = rnd.randn(1, Gs.input_shape[1])
        fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
        images = Gs.run(latents, None, truncation_psi=0.6, \
        randomize_noise=True, output_transform=fmt)
        os.makedirs(config.result_dir, exist_ok=True)
        png_filename = os.path.join(config.result_dir, 'example-'+str(i)+'.png')
        PIL.Image.fromarray(images[0], 'RGB').save(png_filename)

if __name__ == "__main__":
    main()

import os

import pickle

import numpy as np

import PIL.Image

import dnnlib

import dnnlib.tflib as tflib

import config

def main():

tflib.init_tf()

_G, _D, Gs = pickle.load( \

open("results/02051-sgan-faces-2gpu/network-snapshot-021980.pkl", "rb"))

Gs.print_layers()

for i in range(0,1000):

rnd = np.random.RandomState(None)

latents = rnd.randn(1, Gs.input_shape[1])

fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)

images = Gs.run(latents, None, truncation_psi=0.6, \

randomize_noise=True, output_transform=fmt)

os.makedirs(config.result_dir, exist_ok=True)

png_filename = os.path.join(config.result_dir, 'example-'+str(i)+'.png')

PIL.Image.fromarray(images[0], 'RGB').save(png_filename)

if __name__ == "__main__":

main()

Видео

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

cat $(ls ./results/*faces*/fakes*.png | sort --numeric-sort) \
    | ffmpeg -framerate 10 \ # показываем 10 изображений / сек
    -i - # stdin
    -r 25 # выходная частота кадров; кадры будут дублироваться до 25 fps
    -c:v libx264 # x264 для совместимости
    -pix_fmt yuv420p # заставить ffmpeg использовать стандартное
    # цветовое пространство
    -crf 33 # адекватное качество
    -vf "scale=iw/2:ih/2" \ # уменьшаем изображение в 2 раза,
    # т. к. полная детализация не требуется
    -preset veryslow -tune animation
    ./stylegan-facestraining.mp4

cat $(ls ./results/*faces*/fakes*.png | sort --numeric-sort) \

| ffmpeg -framerate 10 \ # показываем 10 изображений / сек

-i - # stdin

-r 25 # выходная частота кадров; кадры будут дублироваться до 25 fps

-c:v libx264 # x264 для совместимости

-pix_fmt yuv420p # заставить ffmpeg использовать стандартное

# цветовое пространство

-crf 33 # адекватное качество

-vf "scale=iw/2:ih/2" \ # уменьшаем изображение в 2 раза,

# т. к. полная детализация не требуется

-preset veryslow -tune animation

./stylegan-facestraining.mp4

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

Можно загрузить готовую модель, обученную на 218794 образцах в течение ~38 GPU-дней. Актуальная версия — от 8 марта 2019 года.

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

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

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

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

Пример переноса обучения StyleGAN на персонажа Holo из “Spice&Wolf”:

64 Ð¾ÑÐ¾Ð±ÑÐ°Ð½Ð½ÑÑ Ð²ÑÑÑÐ½ÑÑ Ð¾Ð±ÑÐ°Ð·ÑÐ° Ð»Ð¸ÑÐ° Holo

И ещё один пример: Луиза из “Zero no Tsukaima”:

Она получилась не такой качественной из-за гораздо меньшего размера выборки.

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

С оригинальной статьёй можно ознакомиться на сайте автора gwern.net.

Редакция

6 мая 2019