[Из песочницы] Генерация изображений из текста с помощью AttnGAN
Привет, Хабр! Представляю вашему вниманию перевод статьи «AttnGAN: Fine-Grained Text to Image Generation with Attentional Generative Adversarial Networks» авторов Tao Xu, Pengchuan Zhang, Qiuyuan Huang, Han Zhang, Zhe Gan, Xiaolei Huang, Xiaodong He.
В данной публикации я хочу рассказать о своих экспериментах с архитектурой AttnGAN для генерации изображений из текстового описания. Эта архитектура уже упоминалась на Хабре после выхода оригинальной статьи в начале 2018 года, и меня заинтересовал вопрос — насколько сложно будет обучить такую модель самостоятельно?
Описание архитектуры
Для тех, кто не знаком с AttnGAN и классическим GAN, — вкратце опишу суть. Классический GAN состоит из минимум 2 нейронных сетей — генератора и дискриминатора. Задачей генератора является порождение некоторых данных (изображений, текста, аудио, видео и т.д), «похожих» на реальные данные из датасета. Задачей дискриминатора является оценка сгенерированных данных, попытка сравнить их с реальными и забраковать. Забракованный результат работы генератора стимулирует его порождать все лучший результат, чтобы «обмануть» дискриминатор, который, в свою очередь, учится все лучше распознавать подделки.
Модификаций GAN существует великое множество, и авторы AttnGAN подошли к вопросу архитектуры весьма изобретательно. Модель состоит из 9 нейросетей, тонко настроенных на взаимодействие. Выглядит это примерно так: 
Энкодеры текста и изображения (text/image encoder) выполняют преобразование исходного текстового описания и реальных изображений в некое внутреннее представление. Характерно, что в данном случае текст рассматривается в виде последовательности отдельных слов, представление которых обрабатывается совместно с представлением изображения, что позволяет сопоставить отдельные слова отдельным частям изображения. Таким образом реализуется механизм внимания, названный авторами статьи DAMSM.
Fca — создает сжатое представление об общей сцене на изображении, исходя из всего текстового описания. Значение C на выходе конкатенируется с вектором из нормального распределения Z, который задает вариативность сцены. Эта информация является основой для работы генератора.
Генератор — самая большая сеть, состоящая из трех уровней. Каждый уровень порождает изображения все большего разрешения, от 64×64 до 256×256 пикселей, и результат работы на каждом уровне корректируется с помощью сетей внимания Fattn, которые несут в себе информацию о правильном расположении отдельных объектов сцены. Кроме того, результаты на каждом уровне проверяются тремя отдельно работающими дискриминаторами, которые оценивают реалистичность изображения и соответствие его общему представлению о сцене.
Обучение
Для тестирования архитектуры я использовал стандартный датасет CUB с фотографиями и текстовым описанием птиц.
Обучение всей модели проходит в два этапа. Первый этап — предобучение сетей DAMSM, состоящей из энкодера текста и изображения. В ходе данного этапа, как описывалось выше, создается «карта внимания», которая выглядит следующим образом:
Как видно из рисунка, DAMSM удается весьма качественно улавливать взаимосвязь между отдельными словами из текстового описания и элементами изображения, особенно легко модели удается распознавать цвета. Надо сказать, что никакой дополнительной информации о том, что такое «красный», «желтый» или «крылья», «клюв» — у системы нет. Есть лишь набор текстов и изображений.
Обучение DAMSM протекает без особых проблем, время обучения на данном датасете 150–200 эпох, что соответствует нескольким часам на GPU высокой мощности.
Второй и основной этап — это обучение генератора с использованием модели DAMSM.
Генератор на каждом уровне порождает изображение все более высокого разрешения — выглядит это так:
Обучение генератора протекает значительно дольше и не всегда так стабильно, рекомендуемое время обучения на данном датасете составляет 300–600 эпох, что соответствует примерно 4–8 суткам на GPU высокой мощности.
Основная проблема при обучении генератора, на мой взгляд, в отсутствии достаточно хороших метрик, которые бы позволяли оценивать качество обучения в более формальном виде. Я изучил несколько реализаций Inception score, который, по идее, позиционируется как универсальная метрика для подобных задач –, но они не показались мне достаточно убедительными. Если вы решили обучить подобный генератор — вам будет необходимо постоянно контролировать ход обучения визуально, по промежуточным результатам. Впрочем, это правило справедливо для любых подобных задач, визуальный контроль всегда необходим.
Результаты
Теперь самое интересное. С помощью обученной модели попробуем сгенерировать изображения, начнем с простых предложений:
Попробуем более сложные описания:
Все текстовые описания являются придуманными, я намеренно не использовал для тестов фразы из датасета. Конечно, не все эти изображения были получены с первой попытки. Модель ошибается, об этом говорят и сами авторы. По мере увеличения текста описания и элементов, подлежащих отображению — становится все сложнее сохранить реалистичность всей сцены. Однако, если вы захотите использовать что-то подобное в продакшне, скажем, генерировать картинки неких объектов для дизайнера, — вы сможете обучить и настроить систему под свои требования, которые могут быть достаточно строгими.
На каждое текстовое описание можно сгенерировать множество вариантов изображений (в том числе и нереалистичных), поэтому всегда будет из чего выбрать.
Технические детали
В данной работе я использовал GPU малой мощности для прототипирования и облачный сервер Google Cloud с установленной Tesla K80 на этапе обучения.
Исходный код был взят из репозитория авторов статьи и прошел серьезный рефакторинг. Система тестировалась на Python 3.6 с Pytorch 0.4.1
Cпасибо за внимание!
Оригинальная статья: AttnGAN: Fine-Grained Text to Image Generation with Attentional Generative Adversarial Networks, 2018 — Tao Xu, Pengchuan Zhang, Qiuyuan Huang, Han Zhang, Zhe Gan, Xiaolei Huang, Xiaodong He.
