3.11.样式迁移

样式迁移

​ 使用卷积神经网络，自动的将一个图像中的风格应用在另一图像之上，即样式迁移(style transfer)

​ 为了完成这一过程，我们需要两张输入图像：一张是内容图像，一张是风格图像，随后使用神经网络修改内容图像，使其在风格上接近风格图像。

1. 基于CNN的样式迁移

​ 奠基性工作

​ 首先初始化合成图像，例如将其初始化为内容图像，该合成图像是样式迁移过程中唯一需要更新的变量，即模型参数。随后选择一个预训练的卷积神经网络来抽取图像的特征，其中的模型参数在训练中无需更新。这个深度卷积神经网络用多个层逐级抽取图像的特征，可以选择其中某些层的输出作为内容特征或风格特征，例如上图中，有3个卷积层，第二层输出内容特征，第一层和第三层输出风格特征。

​ 接下来，通过前向传播（实线）计算损失函数，并通过反向传播(虚线)迭代模型参数，不断更新合成图像，样式迁移常用的损失函数由3部分组成：

1. 内容损失使合成图像与内容图像在内容特征上接近
2. 风格损失使合成图像与风格图像在风格特征上接近
3. 全变分损失则有助于减少合成图像中的噪点

​ 最后，输出结果。

2. 代码实现

2.1 预处理和后处理

'''预处理和后处理'''

rgb_mean = torch.tensor([0.485, 0.456, 0.406])
rgb_std = torch.tensor([0.229, 0.224, 0.225])

#预处理函数对输入图像在RGB3个通道上分别作标准化，并将结果变换成卷积神经网络接受的输入格式
def preprocess(img, image_shape):
    transforms = torchvision.transforms.Compose([
        torchvision.transforms.Resize(image_shape),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(mean=rgb_mean, std=rgb_std)])
    return transforms(img).unsqueeze(0)

#后处理函数将输出图像中的像素值还原回标准化之前的值，由于图像打印函数要求每个像素的浮点数值在0到1之间，则设置为边界
def postprocess(img):
    img = img[0].to(rgb_std.device)
    # 反标准化，乘上标准差加上方差
    img = torch.clamp(img.permute(1, 2, 0) * rgb_std + rgb_mean, 0, 1)
    # ToPILImage 将张量转换成PIL图像格式
    return torchvision.transforms.ToPILImage()(img.permute(2, 0, 1))

2.2 抽取图像特征

pretrained_net = torchvision.models.vgg19(weights=VGG19_Weights.IMAGENET1K_V1)

# vgg19使用了5个卷积块，选择第四卷积快的最后一个卷积层作为内容层
# 因为越往上越全局，所以内容层尽量靠近输出，但样式有局部的也有全局的，所以都均匀取一下
style_layers, content_layers = [0, 5, 10, 19, 28], [25]

# 只要到最高层之间的层数
net = nn.Sequential(*[pretrained_net.features[i] for i in
                      range(max(content_layers + style_layers) + 1)])

# 保留内容层和风格层的输出
def extract_features(X, content_layers, style_layers):
    contents = []
    styles = []
    for i in range(len(net)):
        X = net[i](X)
        if i in style_layers:
            styles.append(X)
        if i in content_layers:
            contents.append(X)
    return contents, styles


'''由于训练时无须改变预训练的VGG的模型参数，所以可以在训练开始之前就提取出内容特征和风格特征'''

def get_contents(image_shape, device):
    content_X = preprocess(content_img, image_shape).to(device)
    contents_Y, _ = extract_features(content_X, content_layers, style_layers)
    return content_X, contents_Y

def get_styles(image_shape, device):
    style_X = preprocess(style_img, image_shape).to(device)
    _, styles_Y = extract_features(style_X, content_layers, style_layers)
    return style_X, styles_Y

2.3 损失函数

​ 由内容损失、风格损失核全变分损失3部分组成

内容损失

​ 内容损失通过平方误差函数衡量合成图像与内容图像在内容特征上的差异，平方误差函数的两个输入均为extract_features函数计算得到的内容层的输出

def content_loss(Y_hat, Y):
    # 我们从动态计算梯度的树中分离目标：
    # 这是一个规定的值，而不是一个变量。
    return torch.square(Y_hat - Y.detach()).mean()

风格损失

​ 为了表达风格层输出的风格，先通过extract_features函数计算风格层的输出。假设该输出的样本数为1，通道数为c，高和宽分别为h和w，我们可以将此输出转换为矩阵$X$,其有$c$行和$hw$列，这个矩阵可以看作c个长度为$hw$的向量$x_1,\cdots ,x_c$组合而成的，其中向量$x_i$代表了通道$i$上的风格特征

​ 使用格拉姆矩阵来表示他们的相关性

在线性代数中，内积空间中一族向量$v_1,\cdots ,v_n$ 的格拉姆矩阵（Gramian matrix、Gram matrix 或 Gramian）是内积的埃尔米特矩阵，其元素由$G_{ij}=<v_i,v_j>$给出。

​ 即格拉姆矩阵中$i$行$j$列的元素$x_{ij}$即向量$x_i$和$x_j$的内积。需要注意的是，当$hw$的值较大时，格拉姆矩阵中的元素容易出现较大的值，此外，格莱姆矩阵的高和宽皆为通道数$c$，为了让风格损失不受这些值的大小影响，下面定义的gram函数将格拉姆矩阵除以了矩阵中元素的个数，即$chw$

def gram(X):#(批量大小，通道数，高，宽)
    num_channels, n = X.shape[1], X.numel() // X.shape[1]
    X = X.reshape((num_channels, n))
    return torch.matmul(X, X.T) / (num_channels * n)
  
def style_loss(Y_hat, gram_Y):
    return torch.square(gram(Y_hat) - gram_Y.detach()).mean()

全变分损失

​ 有时合成图像中有大量高频噪点，即有特别亮或特别暗的颗粒像素，一种常见的去噪方法是全变分去噪(total variation denoising)：

​ 假设$x_{i,j}$表示坐标$(i,j)$处的像素值，降低全变分损失
$$\sum _{i,j}|x_{i,j}-x_{i+1,j}|+|x_{i,j}-x_{i,j+1}|$$
​ 能够尽可能使邻近的像素值相似

def tv_loss(Y_hat):
    return 0.5 * (torch.abs(Y_hat[:, :, 1:, :] - Y_hat[:, :, :-1, :]).mean() +
                  torch.abs(Y_hat[:, :, :, 1:] - Y_hat[:, :, :, :-1]).mean())

损失函数

​ 计算三种损失的加权和，可通过调整权重超参数来调整相对重要性。

content_weight, style_weight, tv_weight = 1, 1e3, 10

def compute_loss(X, contents_Y_hat, styles_Y_hat, contents_Y, styles_Y_gram):
    # 分别计算内容损失、风格损失和全变分损失
    contents_l = [content_loss(Y_hat, Y) * content_weight for Y_hat, Y in zip(
        contents_Y_hat, contents_Y)]
    styles_l = [style_loss(Y_hat, Y) * style_weight for Y_hat, Y in zip(
        styles_Y_hat, styles_Y_gram)]
    tv_l = tv_loss(X) * tv_weight
    # 对所有损失求和
    l = sum(10 * styles_l + contents_l + [tv_l])
    return contents_l, styles_l, tv_l, l

2.4 训练

class SynthesizedImage(nn.Module):
    #将合成的图像视为模型参数，所以前向传播只需返回模型参数即可
    def __init__(self, img_shape, **kwargs):
        super(SynthesizedImage, self).__init__(**kwargs)
        self.weight = nn.Parameter(torch.rand(*img_shape))

    def forward(self):
        return self.weight


def get_inits(X, device, lr, styles_Y):
    gen_img = SynthesizedImage(X.shape).to(device)
    #初始化为图像X
    gen_img.weight.data.copy_(X.data)
    trainer = torch.optim.Adam(gen_img.parameters(), lr=lr)
    #风格图像在各个风格层的格拉姆矩阵styles_Y_gram将在训练前预先计算好。
    styles_Y_gram = [gram(Y) for Y in styles_Y]
    return gen_img(), styles_Y_gram, trainer

'''训练模型'''
def train(X, contents_Y, styles_Y, device, lr, num_epochs, lr_decay_epoch):
    X, styles_Y_gram, trainer = get_inits(X, device, lr, styles_Y)
    scheduler = torch.optim.lr_scheduler.StepLR(trainer, lr_decay_epoch, 0.8)
    animator = d2l.Animator(xlabel='epoch', ylabel='loss',
                            xlim=[10, num_epochs],
                            legend=['content', 'style', 'TV'],
                            ncols=2, figsize=(7, 2.5))
    for epoch in range(num_epochs):
        trainer.zero_grad()
        contents_Y_hat, styles_Y_hat = extract_features(
            X, content_layers, style_layers)
        contents_l, styles_l, tv_l, l = compute_loss(
            X, contents_Y_hat, styles_Y_hat, contents_Y, styles_Y_gram)
        l.backward()
        trainer.step()
        scheduler.step()
        if (epoch + 1) % 10 == 0:
            animator.axes[1].imshow(postprocess(X))
            animator.add(epoch + 1, [float(sum(contents_l)),
                                     float(sum(styles_l)), float(tv_l)])
    return X


device, image_shape = d2l.try_gpu(), (300, 450)
net = net.to(device)
content_X, contents_Y = get_contents(image_shape, device)
_, styles_Y = get_styles(image_shape, device)
output = train(content_X, contents_Y, styles_Y, device, 0.3, 500, 50)
d2l.plt.show()