TensorFlow2实现实时任意风格迁移

人工智能 5月前 0 17

「这是我参加11月更文应战的第19天，活动详情查看：2021最终一次更文应战」

前言

咱们虽然在改善风格搬迁中改善了传统的神经风格搬迁，可是仍然只能运用练习所得的固定数量的风格。因而咱们要学习另一种答应实时恣意风格搬迁的神经网络模型，获得更多构思选择。

自适应实例规范化

AdaIN(adaptive instance normalization) 是实例归一化的一种，这意味着其均值和标准差是在每个图画和每个通道 (H, W) 上核算的。在 CIN 中，和系数是可练习的变量，它们学习不同风格所需的均值和方差。在AdaIN中，和被风格特征的标准差和均值所取代：

AdaIN(x,y)=(y)x−(x)(x)+(y)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” 作为内容层。因而，咱们不需求单独界说内容层。

接下来，咱们需求将对卷积层进行少数但很重要的改善，以改善生成图画的外观。

通过反射填充(reflection padding)削减块伪影

当咱们在卷积层中将填充( 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

解码器结构与完结

尽管编码器运用了4个 VGG 中的网络层( 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 矩阵，以下是风格丢失函数方程式：

Ls=∑i=1L∣∣(i(stylized))−(i(style))∣∣2+∣∣(i(stylized))−(i(style))∣∣2\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

此处， $i\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)