Реализация CycleGAN в Keras и Tensorflow 2.0

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

Весь код из этой статьи будет доступен на моем GitHub здесь.

Таким образом, CycleGAN - это метод перевода изображений. CycleGAN похож на pix2pix, но улучшение заключается в отсутствии необходимых парных наборов данных для обучения. Это означает, что мы можем предоставить CycleGAN любые изображения в качестве ссылок, а любые изображения - в качестве целей стиля, и CycleGAN изучит перевод.

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

pix2pix требует, чтобы выходное изображение было того же ландшафта.

К счастью, многие наборы данных CycleGAN, включая monet2photo, уже доступны для удобного использования в коллекции TensorFlow Datasets (tfds). Если он еще не установлен, вы можете установить его с помощью:

pip3 install tensorflow_datasets

Для этого руководства требуется TensorFlow 2.0, и установка tfds должна автоматически обновить TensorFlow за вас. Если этого не произошло, вы можете обновить до TensorFlow 2.0 вручную, выполнив следующие действия:

pip3 install --upgrade tensorflow

Or:

pip3 install --upgrade tensorflow-gpu

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

train_x и train_y составляют наборы данных X и Y, которые мы обсуждали в предыдущей статье. Первый генератор, G, пытается отобразить train_x в train_y, а второй генератор, F, пытается отобразить train_y в train_x.

После загрузки набора данных мы можем определить несколько гиперпараметров и настроек для модели:

  • LAMBDA, или λ, определяет относительную важность проигрыша цикла по отношению к состязательному проигрышу.
  • img_rows, img_cols и channels представляют размеры изображения, которое мы будем использовать.
  • weight_initializer определяет, какой инициализатор мы будем использовать для установки весов модели.
  • gen_g/f_optimizer и dis_x/y_optimizer определяют оптимизаторы для генераторов и дискриминатора

Затем мы должны установить функцию для нормализации всех изображений до [-1, 1], изменить ее размер в (возможно) пользовательскую форму и изменить ее:

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

Для удобства tfds имеет простой метод dataset.map(function) для отображения функции на весь набор данных. Мы можем использовать это для предварительной обработки всех ранее загруженных наборов данных.

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

Дискриминатор

Дискриминатор имеет один фундаментальный строительный блок: Ck

Напомним из предыдущей статьи: «Ck обозначает слой 4 × 4 Convolution-InstanceNorm-LeakyReLU с k фильтрами и шагом 2»

В коде это может выглядеть примерно так:

Примечание: для использования слоя InstanceNormalization необходимо установить tensorflow_addons, что можно сделать с помощью:

pip3 install tensorflow_addons

Установив Ck, мы можем построить дискриминатор с архитектурой C64, C128, C256, C512.

Как рекомендовано, мы отключаем InstanceNormalization в первом слое Ck и добавляем последний Conv2D слой.

Генератор

Здесь мы будем немного отличаться от того, что рекомендует статья.

Вместо использования модифицированного генератора ResNet в этой статье мы будем использовать модифицированный генератор U-Net, часто встречающийся в таких приложениях, как pix2pix. Генератор U-Net, который мы будем использовать, менее затратен с точки зрения вычислений, и, чтобы сохранить этот учебник как можно более открытым, я решил использовать его вместо генератора ResNet.

Моя реализация генератора ResNet, используемого в статье, доступна на моем GitHub здесь; однако я еще не тестировал и не отлаживал код полностью, поэтому, если вы обнаружите, что с ним что-то не так, сообщите мне.

Dk

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

«dk обозначает 3 × 3 Convolution-InstanceNorm-ReLU с k фильтрами и шагом 2»

Uk

«uk обозначает слой 3 × 3 с дробными полосами-ConvolutionInstanceNorm-ReLU с k фильтрами и шагом 1/2».

Модифицированный U-Net

С блоком понижающей дискретизации Dk и блоком повышающей дискретизации Uk мы можем построить модифицированный генератор U-Net.

U-Net следует стандартной архитектуре кодер-декодер, но особенным генератор делает его «пропускать» соединения. Пропускные соединения подключают части кодировщика напрямую к частям декодера.

Мы начинаем построение архитектуры с определения входа для модели:

Затем мы можем продолжить и создать списки слоев, которые мы будем использовать:

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

Затем мы можем перебрать все слои кодировщика, добавить их в модель и сохранить в skips списке для последующего пропуска подключений.

После добавления всех слоев кодировщика мы должны сделать две вещи:

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

Конечно, первый уровень в декодере не имеет Concatenate() перед ним и, следовательно, не имеет пропуска соединения по причине, о которой мы говорили ранее.

Приведенный выше код оставляет выходные данные модели с формой (?, img_rows/2, img_cols/2, 64). Чтобы придать ему форму (?, img_rows, img_cols, channels), мы должны добавить последний Conv2DTranspose слой с filters, установленным на channels, и вернуть модель:

Вкратце, вот полный код модифицированного генератора U-Net, который мы только что создали:

Затем мы можем быстро создать экземпляры дискриминаторов и генераторов, которые мы только что создали:

Статья TensorFlow pix2pix была отличным помощником при создании этой модели. Я очень рекомендую его всем, кто интересуется pix2pix.

Обучение

Здесь мы отличаемся от классической техники Keras обучения GAN с комбинированными моделями - мы будем использовать объект tf.GradientTape для вычисления потерь.

Объект GradientTape позволяет нам отслеживать операции, выполняемые на графике TensorFlow, и вычислять градиенты относительно некоторых заданных переменных. Я написал руководство, посвященное более подробному объяснению tf.GradientTape, которое доступно здесь.

Убытки

Прежде чем мы сможем настроить функцию поезда, мы должны определить наши потери.

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

Как вы помните, функция потерь дискриминатора математически выглядит примерно так:

Таким образом, эта функция измеряет, насколько близок к 1 вывод дискриминатора для реальных изображений и насколько близок к 0 вывод дискриминатора для поддельных изображений.

В коде вы можете реализовать это примерно так:

Очень важно использовать tf.ones_like и tf.zeros_like вместо np.ones_like и np.zeros_like, поскольку tf.ones/zeros_like возвращают тензоры, а np.ones/zeros_like возвращают массивы numpy. Использование несимвольных тензоров, например массивов numpy, вызовет ошибку, когда GradientTape пытается вычислить градиенты.

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

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

В коде это может выглядеть примерно так:

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

Из-за их сходства мы можем определить одну image_similarity функцию и использовать ее двумя разными способами, чтобы покрыть оба убытка одновременно.

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

tf.reduce_mean(tf.abs(image1-image2)) - это код, эквивалентный принятию ||image1-image2||, средней абсолютной ошибки image1-image2.

Пошаговая функция

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

Мы запускаем функцию со следующего:

  • @tf.function отслеживает все операции в функции и создает вызываемый объект, который может выполнять график TensorFlow.
  • Установка persistent=True в tf.GradientTape означает, что мы можем вызывать tape.gradients несколько раз. По умолчанию после расчета одного набора градиентов лента стирается сама собой, однако, поскольку нам нужно вычислять градиенты несколько раз, persistent должно быть True.

Внутри блока with мы можем начать настройку потерь для дискриминатора Dy.

  • Сначала мы генерируем fake_y, вызывая G на real_x - мы также будем использовать эту переменную позже.
  • Затем мы оцениваем достоверность fake_y, вызывая для него Dy, и сохраняем его в gen_g_validity - эта переменная также будет использоваться позже.
  • Наконец, мы получаем общую потерю Dy, вызывая функцию discriminator_loss, которую мы построили ранее на основе предсказания Dy для real_y и gen_g_validity, которое является предсказанием Dy для fake_y

После расчета потерь для Dy мы можем применить его следующим образом:

Градиенты Dy рассчитываются с использованием tape.gradient. tape.gradient вычисляет градиенты некоторых параметров с учетом потерь. В данном случае потери равны dis_y_loss, а эти параметры равны discriminator_y.trainable_variables.

После этого dis_y_optimizer применяет discriminator_y_gradients ко всем обучаемым переменным Ди.

  • Мы используем tape.stop_recording(), потому что GradientTape не нужно следить за самим вычислением градиентов, что приводит к более быстрым вычислениям.

Мы используем почти идентичный процесс для Dx, с той лишь разницей, что x и y меняются местами, как и G и F

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

Во-первых, настраиваем состязательные проигрыши. Мы делаем это, оценивая, насколько реальными, по мнению дискриминатора, были сгенерированные изображения, принимая gen_loss из gen_g/f_validity.

Затем мы можем настроить потери цикла. Помните, что потеря цикла вычисляет, насколько хорошо переведенное изображение может быть повторно переведено в исходное, G (F (y)) должно ≈ y, а F (G (x)) должно ≈ x.

В приведенном выше коде мы пытаемся преобразовать изображение fake_x/y обратно в real_x/y, а затем измерить потери, используя image_similarity между ними.

Наконец, мы можем настроить потерю идентичности. Как вы помните, потеря идентичности утверждает, что F (x) ≈ x и G (y) ≈ y.

Приведенный выше код является копией того, что мы только что сделали в математике. Он занимает image_similarity между G/F(real_x/y) и real_x/y.

Имея все части потерь генераторов, мы можем сложить все потери в одну переменную:

Член цикла (cyc_x_loss + cyc_y_loss) масштабируется на LAMBDA, чтобы придать ему большее / меньшее значение, чем состязательная потеря. По традиции, потеря идентичности масштабируется на 0.5*LAMBDA для той же цели.

И в тот момент, которого мы все ждали, мы можем, наконец, рассчитать и применить градиенты для генераторов, используя тот же метод, который мы использовали раньше:

И это полноценная функция шага поезда!

Напомним, вот весь код, который мы только что реализовали:

Цикл обучения

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

Для каждой эпохи прокрутите каждый пример данных, измените его форму на (1, img_rows, img_cols, channels) и передайте в функцию step.

Если вы хотите визуализировать прогресс CycleGAN, вы можете реализовать какую-то функцию вроде:

И тогда мы можем назвать это каждую эпоху или около того:

Итак, мы идем!

Оставьте это работать примерно на 45 эпох, и вы должны вернуться к чему-то вроде:

Чтобы сделать еще один шаг вперед, вы можете разделить видео на отдельные кадры, обрезать и изменить их размер до 256 на 256 изображений, преобразовать каждый кадр в фотографию Моне и сшить результат обратно в видео. Возможности безграничны!

Заключение

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

Я бы порекомендовал поиграть с реализацией CycleGAN, изменив функцию оптимизации или архитектуры, чтобы увидеть, сможете ли вы найти что-то еще более оптимальное. Вы также можете создать свою версию модифицированного генератора ResNet или использовать ту, что есть на моем GitHub.

Как всегда, готовый код для этой статьи доступен в форматах .ipynb и .py на моем GitHub здесь.

Если у вас есть какие-либо вопросы, не стесняйтесь оставлять их в ответ на эту статью, и я постараюсь ответить вам как можно скорее.

И до следующего раза,
Удачного кодирования!

Дальнейшее чтение

Бумага CycleGAN: https://arxiv.org/pdf/1703.10593.pdf

pix2pix: https://arxiv.org/abs/1611.0700
U-Net: https://arxiv.org/abs/1505.04597

Блокнот TensorFlow на CycleGAN (очень помог в создании этой статьи): https://www.tensorflow.org/tutorials/generative/cyclegan

Моя предыдущая статья о CycleGAN: https://medium.com/analytics-vidhya/the-beauty-of-cyclegan-c51c153493b8
Моя предыдущая статья о GAN: https://medium.com/analytics- vidhya / implementation-a-gan-in-keras-d6c36bc6ab5f