Tensorflow2 реализует WGAN и WGAN-GP

«Это 12-й день моего участия в ноябрьском испытании обновлений, ознакомьтесь с подробностями события:Вызов последнего обновления 2021 г."

Построение WGAN (Wasserstein GAN)

С момента появления GAN во многих документах предпринимались попытки решить проблему нестабильности обучения GAN с помощью эвристики, такой как тестирование различных сетевых архитектур, гиперпараметров и оптимизаторов. С введением Wasserstein GAN (WGAN) исследования по этой проблеме достигли крупного прорыва.

WGAN облегчает или даже устраняет многие проблемы в процессе обучения GAN. Фундаментальным улучшением по сравнению с оригинальным GAN является модификация функции потерь. Теоретически, если два распределения не пересекаются, расхождение JS больше не будет непрерывным и, следовательно, будет недифференцируемым, что приведет к нулевому градиенту. WGAN решает эту проблему, используя новую функцию потерь, которая непрерывна и везде дифференцируема!

Описание потери Вассерштейна

Мы все знакомы с целевой функцией оригинальной GAN, поэтому вот краткий обзор:

в,$D$представляет дискриминатор,$G$представляет генератор,$x$представляет реальные данные,$z$представляет скрытую переменную. Преобразовав приведенную выше форму, можно получить следующую форму функции ценности:

WGAN использует новую функцию потерь, называемую расстоянием бульдозера или расстоянием Вассерштейна. Он используется для измерения расстояния или усилия, необходимого для преобразования одного распределения в другое. Математически это минимальное расстояние для каждого совместного распределения между реальным изображением и сгенерированным изображением, и функция значения WGAN принимает вид:

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

Это среднее значение выходного сигнала дискриминатора, умноженное на -1. мы используем$y_i$В виде меток, где +1 представляет собой реальное изображение, а -1 представляет поддельное изображение. Следовательно, мы можем реализовать потерю Вассерштейна как пользовательскую функцию потерь TensorFlow Keras следующим образом:

def wasserstein_loss(self, y_true, y_pred):
	w_loss = -tf.reduce_mean(y_true*y_pred)
	return w_loss

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

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

1- Реализация ограничения Липшица

Математическое предположение, упомянутое в потере Вассерштейна, представляет собой 1-липшицеву функцию. Назовем критику D(x) 1-липшицевой, если она удовлетворяет следующему неравенству:

для двух изображений$x_1$и$x_2$, абсолютное значение выходной разницы критика должно быть меньше или равно абсолютному значению его средней попиксельной разницы. Другими словами, вывод критика не должен сильно различаться для разных изображений, будь то настоящие или фальшивые. Когда была предложена WGAN, авторы не могли придумать подходящей реализации для реализации этого неравенства. Поэтому они придумали способ снизить вес критиков до какой-то небольшой величины. Таким образом, вывод слоя и, в конечном счете, вывод критика ограничивается небольшим значением. В документе WGAN веса ограничены диапазоном [-0,01, 0,01]. Отсечение веса может быть реализовано двумя способами. Один из способов — написать пользовательскую функцию ограничения и использовать ее при создании нового слоя следующим образом:

class WeightsClip(tf.keras.constraints.Constraint):
    def __init__(self, min_value=-0.01, max_value=0.01):
        self.min_value = min_value
        self.max_value = max_value
    def __call__(self, w):
        return tf.clip_by_value(w, self.min, self.max_value)

Затем функцию можно передать слою, который принимает функцию ограничения, следующим образом:

model = tf.keras.Sequential(name='critics')        
model.add(Conv2D(16, 3, strides=2, padding='same',
				kernel_constraint=WeightsClip(),
				bias_constraint=WeightsClip()))
model.add(BatchNormalization(
beta_constraint=WeightsClip(),
gamma_constraint=WeightsClip()))

Однако добавление кода ограничений во время создания каждого слоя может привести к раздуванию кода. Поскольку нам не нужно выбирать слои для обрезки, мы можем использовать цикл для считывания весов и записи их обратно после обрезки, например так: Для слоев в comment.layers:

for layer in critic.layers:
    weights = layer.get_weights()
    weights = [tf.clip_by_value(w, -0.01, 0.01) for w in weights]
    layer.set_weights(weights)

тренировочный процесс

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

for _ in range(self.n_critic):
    real_images = next(data_generator)
    critic_loss = self.train_critic(real_images,     batch_size)

Этапы обучения генератора:

self.critic = self.build_critic()
self.critic.trainable = False
self.generator = self.build_generator()
critic_output = self.critic(self.generator.output)
self.model = Model(self.generator.input, critic_output)
self.model.compile(loss = self.wasserstein_loss, optimizer = RMSprop(3e-4))
self.critic.trainable = True

В предыдущем коде, установивtrainable = FalseЗаморозьте слой комментатора и свяжите его с генератором, чтобы создать новую модель и скомпилировать ее. После этого мы можем сделать критика обучаемым, что не повлияет на нашу уже скомпилированную модель.

Мы используемtrain_on_batch()API выполняет один этап обучения, который автоматизирует прямые вычисления, вычисления потерь, обратное распространение и обновления весов:

g_loss = self.model.train_on_batch(g_input,   real_labels)

На следующем рисунке показана архитектура генератора WGAN:

На следующем рисунке показана критически важная архитектура WGAN:

Несмотря на улучшение по сравнению с исходной GAN, обучение WGAN затруднено, а качество получаемых изображений не лучше, чем у исходной GAN. Далее будет реализован вариант WGAN, WGAN-GP, который быстрее обучается и дает более четкие изображения.

Реализовать штраф за градиент (WGAN-GP)

Как признают авторы WGAN, отсечение веса не является идеальным способом соблюдения липшицевых ограничений. У него есть два недостатка: недостаточное использование пропускной способности сети и взрывные/исчезающие градиенты. Когда мы урезаем веса, мы также ограничиваем обучаемость критиков. Отсечение веса заставляет сеть изучать только простые функции. Следовательно, возможности нейронной сети становятся недоиспользованными. Во-вторых, значение урожая должно быть тщательно отрегулировано. Если установить слишком высокое значение, градиент взорвется, нарушив ограничение Липшица. Если установлено слишком низкое значение, градиент исчезнет при обратном распространении сети. Точно так же отсечение веса толкает градиент к двум крайностям, как показано на следующем рисунке:

Таким образом, Gradient Penalty (GP) предлагается заменить отсечение веса для обеспечения ограничений Липшица следующим образом:

Мы рассмотрим каждую переменную в уравнении и реализуем их в коде.

мы обычно используем$x$представляет реальное изображение, но теперь уравнение имеет$\hat x$.$\hat x$представляет собой точечную интерполяцию между реальным изображением и поддельным изображением. Отношение изображения (эпсилон) получено из равномерного распределения [0,1]:

epsilon = tf.random.uniform((batch_size,1,1,1))
interpolates = epsilon*real_images + (1-epsilon)*fake_images

в соответствии сДокумент ВГАН-ГП, для наших целей мы можем понимать это так, поскольку градиенты получаются из смеси реальных и поддельных изображений, поэтому нам не нужно вычислять потери для настоящих и поддельных изображений отдельно.

$\nabla _{\hat x}D(\hat x)$term - это градиент критического выхода по отношению к интерполяции. мы можем использовать сноваtf.GradientTape()чтобы получить градиент:

with tf.GradientTape() as gradient_tape:
	gradient_tape.watch(interpolates)
	critic_interpolates = self.critic(interpolates)
	gradient_d = gradient_tape.gradient(critic_interpolates, [interpolates])

Следующим шагом является расчет нормы L2:

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

grad_loss = tf.square(grad)
grad_loss = tf.reduce_sum(grad_loss, axis=np.arange(1, len(grad_loss.shape)))
grad_loss = tf.sqrt(grad_loss)

в исполненииtf.reduce_sum(), мы исключаем первое измерение на оси, поскольку это измерение является размером партии. Штраф направлен на приближение нормы градиента к 1, что является последним шагом в вычислении потери градиента:

grad_loss = tf.reduce_mean(tf.square(grad_loss - 1))

в уравнении$λ$это отношение штрафа за градиент к другим критическим потерям, установленное здесь на 10. Теперь мы добавляем все критические потери и штрафы за градиент к обратному распространению и обновляем веса:

total_loss = loss_real + loss_fake + LAMBDA * grad_loss
gradients = total_tape.gradient(total_loss, self.critic.variables)
self.optimizer_critic.apply_gradients(zip(gradients, self.critic.variables))

Это все, что нужно добавить в WGAN, чтобы сделать его WGAN-GP. Однако необходимо удалить следующие разделы:

1. отсечение веса
2. Пакетная нормализация среди критиков

Штраф за градиент заключается в наложении штрафа на норму градиента рецензента независимо для каждого ввода. Однако нормализация партии изменяет градиенты со статистикой партии. Чтобы избежать этой проблемы, пакетная нормализация убрана из критики. Архитектура критика такая же, как WGAN, но не включает нормализацию пакетов:

Вот образцы, сгенерированные обученным WGAN-GP:

Они выглядят четкими и красивыми, очень похожими на образцы в наборе данных Fashion-MNIST. Обучение очень стабильное и сходится очень быстро!