Реализация 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
списке для последующего пропуска подключений.
После добавления всех слоев кодировщика мы должны сделать две вещи:
- Подготовьте список пропусков для подключения, перевернув его и удалив первый элемент. Мы должны перевернуть список, чтобы он был в правильном порядке, и удалить первый элемент, потому что он подключен прямо к декодеру - здесь нет необходимости пропускать подключение.
- Добавьте все слои декодера; мы используем слой
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