TensorFlow2 реализует передачу произвольного стиля в реальном времени.

«Это 19-й день моего участия в ноябрьском испытании обновлений. Подробную информацию об этом событии см.:Вызов последнего обновления 2021 г."

предисловие

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

Адаптивная нормализация экземпляра

AdaIN(adaptive instance normalization)является типом нормализации экземпляра, что означает, что его среднее значение и стандартное отклонение(H, W)рассчитано выше. существуетCINсередина, $γ$ и $бета$ Коэффициенты — это обучаемые переменные, которые изучают среднее значение и дисперсию, необходимые для разных стилей. В АдаИН, $γ$ и $бета$ заменены стандартным отклонением и средним значением характеристик стиля:

AdaIN(x,y)=\sigma(y)\frac {x-\mu (x)}{\sigma(x)} + \mu(y)

AdaINВсе еще можно понимать как форму условной нормализации экземпляра, где условия являются функциями стиля, а не метками стиля. Во время обучения и логического вывода мы используем VGG для извлечения выходных данных слоя стилей и используем его статистику в качестве условий стиля, что избавляет от необходимости предварительно определять только фиксированный набор стилей. Используйте TensorFlow для создания пользовательского слоя AdaIN:

class AdaIN(layers.Layer):
    def __init__(self, epsilon=1e-5):
        super(AdaIN, self).__init__()
        self.epsilon = epsilon
    
    def call(self, inputs):
        x = inputs[0] # content
        y = inputs[1] # style
        mean_x, var_x = tf.nn.moments(x, axes=(1,2), keepdims=True)
        mean_y, var_y = tf.nn.moments(y, axes=(1,2), keepdims=True)
        std_x = tf.sqrt(var_x + self.epsilon)
        std_y = tf.sqrt(var_y + self.epsilon)
        output = std_y * (x - mean_x) / std_x + mean_y
        return output

Tips：可以看出，这是对 AdaIN 方程式的直接实现。其中 tf.nn.moments 用于计算特征图的均值和方差，其中轴1、2指向特征图的H，W。还设置keepdims = True以使结果保持四个维度，形状为(N, 1, 1, C)，而不是默认值(N, C)。前者允许TensorFlow使用形状为(N, H, W, C)的输入张量使用广播机制。

Далее мы будемAdaINИнтегрируйте в передачу стиля.

сеть передачи стилей

На следующем рисунке показана архитектура и процесс обучения сети передачи стилей:

风格迁移网络的架构和训练流程

сеть передачи стилей (style transfer network, STN)Да编码器/解码器сеть, где кодировщик использует фиксированный VGG для кодирования содержимого и функций стиля. Затем AdaIN кодирует стилистические особенности в статистику характеристик контента, а декодер использует эти новые функции для создания стилизованных изображений.

Структура кодировщика и реализация

Кодер используетVGGСборка получает:

    def build_encoder(self, name='encoder'):
        self.encoder_layers = [
            'block1_conv1',
            'block2_conv1', 
            'block3_conv1', 
            'block4_conv1']
        vgg = tf.keras.applications.VGG19(include_top=False,weights='imagenet')

        layer_outputs = [vgg.get_layer(x).output for x in self.encoder_layers]

        return Model(vgg.input, layer_outputs, name=name)

Tips：这类似于神经风格迁移，使用最后一个风格层 ”block4_conv1” 作为内容层。因此，我们不需要单独定义内容层。

Далее нам нужно внести небольшие, но существенные улучшения в сверточные слои, чтобы улучшить внешний вид результирующего изображения.

Уменьшите блокирующие артефакты с помощью заполнения отражения

Когда мы помещаем отступы в сверточный слой (padding ) при применении к входному тензору дополняет тензор постоянными нулями. Однако внезапное падение значения на границе создает высокочастотные компоненты и вызывает блокирующие артефакты в результирующем изображении. Одним из способов уменьшения этих высокочастотных составляющих является добавление полных вариационных потерь (total variation loss) как регуляризатор:

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

Суммарные вариационные потери равны $L_1$ Сумма норм. Таким образом, обучение будет пытаться минимизировать эту функцию потерь, чтобы уменьшить высокочастотные компоненты.

Другой вариант — заменить постоянные нули в заполнении отраженным значением. Например, если мы заполним нулями[10, 8, 9]массив, вы получите[0, 10, 8, 9, 0], мы видим, что значение между 0 и его соседями меняется довольно резко.

Если мы используем заполнение отражения, массив заполнения будет[8, 10, 8, 9, 8], что обеспечит более плавный переход к границе. но,Keras Conv2DОтражающая прокладка не поддерживается, поэтому нам нужно использоватьTensorFlowСоздать пользовательскийConv2D. В следующем фрагменте кода показано добавление отражающего заполнения к входному тензору перед сверткой:

class Conv2D(layers.Layer):
    def __init__(self, in_channels, out_channels, kernel=3, use_relu=True):
        super(Conv2D, self).__init__()
        self.kernel = kernel
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.use_relu = use_relu
    
    def build(self, input_shape):
        self.w = self.add_weight(shape=[
		            self.kernel,
		            self.kernel,
		            self.in_channels,
		            self.out_channels],
		            initializer='glorot_normal',
		            trainable=True, name='bias')
        
        self.b = self.add_weight(shape=(
		            self.out_channels,),
		            initializer='zeros',
		            trainable=True,
		            name='bias')

    @tf.function
    def call(self, inputs):
        padded = tf.pad(inputs, [[0,0],[1,1],[1,1],[0,0]], mode='REFLECT')
        # perform conv2d using low level API
        output = tf.nn.conv2d(padded, self.w, strides=1, padding='VALID') + self.b
        if self.use_relu:
            output = tf.nn.relu(output)
        return output

Структура и реализация декодера

Хотя кодировщик использует 4VGGСетевой уровень в (block1_conv1прибытьblock4_conv1), но AdaIN использует только последний уровень кодировщикаblock4_conv1. Следовательно, входной тензор декодера такой же, какblock4_conv1Вывод слоя активации такой же. Декодер состоит из сверточных слоев и слоев повышающей дискретизации, как показано в следующем коде:

    def build_decoder(self):
        block = tf.keras.Sequential([
            Conv2D(512, 256, 3),
            UpSampling2D((2,2)),
            Conv2D(256,256,3),
            Conv2D(256,256,3),
            Conv2D(256,256,3),
            Conv2D(256,128,3),
            UpSampling2D((2,2)),
            Conv2D(128,128,3),
            Conv2D(128,64,3),
            UpSampling2D((2,2)),
            Conv2D(64,64,3),
            Conv2D(64,3,3,use_relu=False)
        ], name='decoder')
        return block

Tips：前面的代码使用具有反射填充的自定义Conv2D。除不具有任何非线性激活函数的输出层外，所有层均使用ReLU激活函数。

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

Предварительная обработка VGG

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

    def preprocess(self, image):
        # RGB to BGR
        image = tf.reverse(image, axis=[-1])
        return tf.keras.applications.vgg19.preprocess_input(image)

Мы можем сделать обратное в постобработке, то есть добавить среднее значение цвета и инвертировать цветовые каналы. Однако это то, что может узнать декодер, поскольку среднее значение цвета эквивалентно смещению в выходном слое. Таким образом, мы позволим процессу обучения выполнять постобработку, и все, что мы будем делать, это обрезать пиксели до[0,255]Сфера:

    def postprocess(self, image):
        return tf.clip_by_value(image, 0., 255.)

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

Реализация сети передачи стилей

структураSTNочень просто, просто подключите编码器,AdaINи解码器Вот и все, как показано на предыдущей схеме архитектуры.STNИли модель, которую мы будем использовать для выполнения логического вывода. Код для этого выглядит следующим образом:

        """
        Style Transfer Network
        """
        content_image = self.preprocess(content_image_input)
        style_image = self.preprocess(style_image_input)

        self.content_target = self.encoder(content_image)
        self.style_target = self.encoder(style_image)

        adain_output = AdaIN()([self.content_target[-1], self.style_target[-1]])
        
        self.stylized_image = self.postprocess(self.decoder(adain_output))

        self.stn = Model([content_image_input, style_image_input], self.stylized_image)

Изображения содержимого и стилей предварительно обрабатываются, а затем передаются в кодировщик. последний векторный слойblock4_conv1ВходитьAdaIN(). Затем стилизованные элементы передаются в декодер для создания изображений, стилизованных под RGB.

Обучение модели переноса произвольного стиля в реальном времени

Подобно нейронной передаче стиля, потеря содержания и потеря стиля основаны на фиксированныхVGGИзвлеченные активации вычисляются. Потеря контента также $L_2$ норму, а теперь сравните содержательные особенности получившихся стилизованных изображений сAdaIN, а не функции изображения содержимого, как показано в следующем коде, что приводит к более быстрой сходимости:

            content_loss = tf.reduce_sum((output_features[-1]-adain_output)**2)

Для потери стиля обычно используетсяGramМатрица заменена средним значением и активирована дисперсия $L_2$ норма. Это дает результаты, аналогичныеGramматрица, следующее уравнение функции потери стиля:

\mathcal L_s = \sum_{i=1}^L||\mu( \phi_i(stylized) )-\mu(\phi_i(style))||_2+||\sigma( \phi_i(stylized) )-\sigma(\phi_i(style))||_2

здесь, $\phi_i$ выражатьVGG-19Слой, используемый для вычисления потери стиля в файлах .

мы вAdaINиспользуется в слояхtf.nn.momentsдля вычисления статистической суммы между функциями из стилизованного изображения и изображения стиля $L_2$ норма, мы усредняем потери слоя контента следующим образом:

    def calc_style_loss(self, y_true, y_pred):
        n_features = len(y_true)
        epsilon = 1e-5
        loss = []

        for i in range(n_features):
            mean_true, var_true = tf.nn.moments(y_true[i], axes=(1,2), keepdims=True)
            mean_pred, var_pred = tf.nn.moments(y_pred[i], axes=(1,2), keepdims=True)
            std_true = tf.sqrt(var_true + epsilon)
            std_pred = tf.sqrt(var_pred + epsilon)
            mean_loss = tf.reduce_sum(tf.square(mean_true-mean_pred))
            std_loss = tf.reduce_sum(tf.square(std_true-std_pred))
            loss.append(mean_loss + std_loss)
        
        return tf.reduce_mean(loss)

Последним шагом является запись шага обучения:

    @tf.function
    def train_step(self, train_data):
        with tf.GradientTape() as tape:
            adain_output, output_features, style_target = self.training_model(train_data)

            content_loss = tf.reduce_sum((output_features[-1]-adain_output)**2)
            style_loss = self.style_weight * self.calc_style_loss(style_target, output_features)
            loss = content_loss + style_loss

            gradients = tape.gradient(loss, self.decoder.trainable_variables)

            self.optimizer.apply_gradients(zip(gradients, self.decoder.trainable_variables))
        
        return content_loss, style_loss

	def train(self, train_generator, test_generator, steps, interval=500, style_weight=1e4):
        self.style_weight = style_weight

        for i in range(steps):
            train_data = next(train_generator)
            content_loss, style_loss = self.train_step(train_data)

            if i % interval == 0:
                ckpt_save_path = self.ckpt_manager.save()
                print ('Saving checkpoint for step {} at {}'.format(i, ckpt_save_path))
                print(f"Content_loss {content_loss:.4f}, style_loss {style_loss:.4f}")
                val_data = next(test_generator)
                self.stylized_images = self.stn(val_data)
                self.plot_images(val_data[0], val_data[1], self.stylized_images)

Tips：我们将内容权重固定为1，并调整风格权重，在示例中，我们将风格权重设置为1e4。在上面展示的网络架构图中，看起来好像有三个要训练的网络，但是其中两个是冻结参数的VGG，因此唯一可训练的网络是解码器。

Отображение большего эффекта

В обучающем примере используйте лица в качестве изображений содержимого и используйтеcyclegan/monet2photoкак стиль изображения. Хотя картины Моне принадлежат к одному художественному стилю, с точки зрения стилевой передачи каждое стилевое изображение является уникальным стилем.monet2photoНабор данных содержит1193стиль изображения, что означает, что мы будем использовать1193различные стили для обучения сети! На изображении ниже показан пример изображения, сгенерированного нашей сетью:

效果展示

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

последовательное соединение

TensorFlow2 реализует передачу нейронного стиля

Улучшение передачи нейронного стиля