TensorFlow2实现实时任意风格迁移

2021-11-22 09:32:47 浏览数 (1)

前言

我们虽然在改进风格迁移中改进了传统的神经风格迁移,但是仍然只能使用训练所得的固定数量的风格。因此我们要学习另一种允许实时任意风格迁移的神经网络模型,获得更多创意选择。

自适应实例规范化

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(x,y)=σ(y)σ(x)x−μ(x)​ μ(y)

AdaIN 仍可以理解为条件实例规范化的一种形式,其中条件是风格特征而不是风格标签。在训练和推理时,我们使用VGG提取风格层输出并将其统计信息用作风格条件,这样避免了只能预先定义一组固定风格。 使用TensorFlow来创建自定义AdaIN层:

代码语言:javascript复制
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 构建获得的:

代码语言:javascript复制
    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 )作为正则化器:

  1. 首先,通过将图像移动一个像素来计算高频分量,
  2. 然后减去原始图像以创建一个矩阵。

总变分损失是 L1L_1L1​ 范数的总和。因此,训练将尝试最小化此损失函数,以减少高频分量。

还有另一种选择,就是用反射值替换填充中的常数零。例如,如果我们用零填充 [10, 8, 9] 的数组,则将得到 [0, 10, 8, 9, 0] ,我们可以看到0和它的邻居之间的值的变化十分突然。

如果我们使用反射填充,则填充数组将为 [8, 10, 8, 9, 8] ,这将向边界提供更平滑的过渡。但是,Keras Conv2D 不支持反射填充,因此我们需要使用 TensorFlow 创建自定义 Conv2D 。以下代码片段显示了在卷积之前向输入张量中添加反射填充:

代码语言:javascript复制
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_conv1block4_conv1 ),但AdaIN仅使用编码器的最后一层 block4_conv1 。因此,解码器的输入张量与 block4_conv1 的激活层输出相同。解码器由卷积和上采样层组成,如以下代码所示:

代码语言:javascript复制
    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然后减去颜色均值来对图像进行预处理,代码如下:

代码语言:javascript复制
    def preprocess(self, image):
        # RGB to BGR
        image = tf.reverse(image, axis=[-1])
        return tf.keras.applications.vgg19.preprocess_input(image)
复制代码

我们可以在后期处理中进行反向操作,即添加颜色均值并反转颜色通道。但是,这是解码器可能会学到的,因为颜色均值等效于输出层中的偏置。因此,我们将让训练过程来完成后期处理工作,而我们要做的仅仅是将像素裁剪至 [0,255] 范围:

代码语言:javascript复制
    def postprocess(self, image):
        return tf.clip_by_value(image, 0., 255.)
复制代码

现在,我们已经准备好所有构件,剩下要做的就是将它们放在一起以创建 STN 和训练过程。

实现风格迁移网络

构造 STN 非常简单,只需连接编码器AdaIN解码器即可,如前面的架构图所示。 STN 还是我们将用来执行推理的模型。执行此操作的代码如下:

代码语言:javascript复制
        """
        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 提取的激活来计算的。内容损失也是 L2L_2L2​ 范数,但是现在将生成的风格化图像的内容特征与 AdaIN 的输出进行比较,而不是与内容图像中的特征进行比较,如以下代码所示,这使收敛速度更快:

代码语言:javascript复制
            content_loss = tf.reduce_sum((output_features[-1]-adain_output)**2)
复制代码

对于风格损失,将常用的 Gram 矩阵替换为均值和方差激活的 L2L_2L2​ 范数。这产生的结果类似于 Gram 矩阵,以下是风格损失函数方程式:

Ls=∑i=1L∣∣μ(ϕi(stylized))−μ(ϕi(style))∣∣2 ∣∣σ(ϕi(stylized))−σ(ϕi(style))∣∣2mathcal L_s = sum_{i=1}^L||mu( phi_i(stylized) )-mu(phi_i(style))||_2 ||sigma( phi_i(stylized) )-sigma(phi_i(style))||_2Ls​=i=1∑L​∣∣μ(ϕi​(stylized))−μ(ϕi​(style))∣∣2​ ∣∣σ(ϕi​(stylized))−σ(ϕi​(style))∣∣2​

此处,ϕiphi_iϕi​ 表示 VGG-19 中用于计算风格损失的层。

我们在 AdaIN 层中使用 tf.nn.moments 来计算来自风格化图像和风格图像的特征之间的统计量和 L2L_2L2​ 范数,我们对内容层的损失求均值,如下所示:

代码语言:javascript复制
    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)
复制代码

最后一步是编写训练步骤:

代码语言:javascript复制
    @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 种不同样式来训练网络!下图显示了由我们的网络生成的图像示例:

上图中的显示了使用网络训练时未使用的风格图像(即测试风格数据集)在推理时进行风格迁移的情况。每种风格迁移仅通过单次前向计算进行,这比原始神经风格迁移算法的迭代优化快得多。

0 人点赞