25PyTorch 实现图像风格迁移

图像风格迁移

图像的风格迁移其实就是利用相关算法对一些著名画作的风格进行学习,然后再把这种风格应用到我们熟悉的图片中。该技术最早由 Gatys 等人 提出,并且将算法应用于他们所发布的软件 Prisma 中。由于该技术不像传统的图像处理软件一般,直接对像素进行操作,而是采用神经网络相关算法模拟名家的绘画风格。因此,在软件发布之初,就吸引了上千万的融资。本实验将对该篇论文中的风格迁移技术进行详细的讲解,并且利用 Pytorch 对其进行实现。
下图为本实验最终的结果展示图:

image.png

从上图可以看出,我们模型的数据由两张原始图片(即内容图像和风格图像)和一张合成的新图像组成。 其中内容图像是一个可爱的小孩纸,而风格图像是具有浓浓的和式水墨风格的图像。

数据的加载

在一般的神经网络课程中,我们需要大量的数据来保证模型的泛化性。而在图像的风格迁移中,我们只需要两张图片(即内容图像和风格图像)即可。让我们来加载这两张图片:

!wget -nc "https://labfile.oss.aliyuncs.com/courses/861/test_pics.zip"
!unzip -o "test_pics.zip"

让我们利用 OpenCV 来对这两张图片进行展示:

import cv2
import matplotlib.pyplot as plt  # plt 用于显示图片
%matplotlib inline

content_path = "content.png"
style_path = "style.png"
plt.subplot(121)  # 1行两列,第一个
figure = cv2.imread(content_path)
# 这里需要指定利用 cv 的调色板,否则 plt 展示出来会有色差
plt.imshow(cv2.cvtColor(figure, cv2.COLOR_BGR2RGB))

plt.subplot(122)  # 1行两列,第二个
figure = cv2.imread(style_path)
# 这里需要指定利用 cv 的调色板,否则 plt 展示出来会有色差
plt.imshow(cv2.cvtColor(figure, cv2.COLOR_BGR2RGB))

由于这两张图片的原始大小不同,且为了保证后面放入任何图片都可以对其进行迁移。我们需要对图片进行预处理操作:
将图片大小缩放为512×512。
将其类型转为 Tensor。
代码如下:

import PIL.Image as Image
import torchvision.transforms as transforms
img_size = 512


def load_img(img_path):
    img = Image.open(img_path).convert('RGB')
    img = img.resize((img_size, img_size))
    img = transforms.ToTensor()(img)
    # 为img增加一个维度:1
    # 因为神经网络的输入为 4 维
    img = img.unsqueeze(0)
    return img

让我们来测试一下上面的函数。加载这两张图片,并将其处理成神经网络能够使用的类型:

import torch
from torch.autograd import Variable
# 判断环境是否支持GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 加载风格图片
style_img = load_img(style_path)
# 对img进行转换为 Variable 对象,使它能够动态计算梯度
style_img = Variable(style_img).to(device)
# 加载内容图片
content_img = load_img(content_path)
content_img = Variable(content_img).to(device)
print(style_img.size(), content_img.size())

图像风格迁移算法

该算法主要使用三张图片,一张输入图片 G ,一张内容图片 C 和一张风格图片 S。

为了衡量任意两张图片的差距,我们还需要定义两个函数式,如下:
计算两张图片内容之间的差距 D_C(即内容损失函数)。
计算两张图片风格的差距 D_S(即风格损失函数)。

图像风格迁移的核心思想:输入图片 G,并且改变这张图片。使输入图片 G 与内容图片 C 之间的内容间距 D_C 最小,进而达到新图片 GG 与内容图片 CC 的内容一致的目的。使输入图片 G 与风格图片 S 之间的风格间距 D_S 最小,进而达到图片 G 与图片 S 风格一致的目的。
或许,你对上面的核心思想你还有很多的疑问,对模型的输入输出还有很多的不解(这也是整个风格迁移算法中最有意思的一部分)。但是,我们需要两个损失函数来衡量图片的内容差距和风格差距,这一定论并不会有什么争议。
因此,让我们先来对这两个损失函数进行讲解,之后再来阐述具体的神经网络模型和训练步骤。

各个损失的计算步骤如下所示:


image.png

损失函数

模型的总损失由内容损失和风格损失加权而成,定义总损失如下:

image.png

其中两个损失的权重α和β可以自行设定。不同的设定方式表示了你对内容和风格的看重程度。
内容损失函数
我们的内容损失函数采用的是最传统的交叉熵损失:
image.png

image.png

由于 Pytorch 为我们提供了完整的交叉熵损失函数。因此,我们能够很容易的定义内容损失函数,如下:

import torch.nn as nn


class Content_Loss(nn.Module):
    # 其中 target 表示 C ,input 表示 G,weight 表示 alpha 的平方
    def __init__(self, target, weight):
        super(Content_Loss, self).__init__()
        self.weight = weight
        # detach 可以理解为使 target 能够动态计算梯度
        # target 表示目标内容,即想变成的内容
        self.target = target.detach() * self.weight
        self.criterion = nn.MSELoss()

    def forward(self, input):
        self.loss = self.criterion(input * self.weight, self.target)
        out = input.clone()
        return out

    def backward(self, retain_variabels=True):
        self.loss.backward(retain_graph=retain_variabels)
        return self.loss

上述代码其实就是简单的交叉熵损失。而其中的 weightweight 参数其实就是将公式中的α放到了括号里面,即:


image.png

为了测试上面代码,让我们来构造一个以 content_img 为目标的损失函数,该函数可以计算任意一张图像与目标 content_img 之间的内容损失:

# 损失函数的测试
cl = Content_Loss(content_img, 1)
# 随机图片
rand_img = torch.randn(content_img.data.size(), device=device)
cl.forward(rand_img)
print(cl.loss)

风格损失函数
风格是一种很难说清楚的概念,假设我们通过一个神经网络模型,对图片的特征进行了提取。那么我们又应该如何比较这些特征之间的风格呢?
在计算风格损失之前,我们首先需要提取图片的风格,然后再利用一些常用的损失函数,计算两种风格之间的差距与损失。那么我们应该怎样去获得图片的风格呢?
对于风格的提取,我们需要用到 Gram 矩阵。Gram 矩阵由i通道的特征图与j通道的特征图的内积计算而成。这个值可以表示为 ii 通道的特征图 与 jj 通道的特征图的互相关程度。
如果特征图F的通道数为n,则计算得到的 Gram 矩阵的大小则为n×n。且该矩阵的第i行第j列的值可以表示为第i个特征图和第j个特征图之间的互相关程度。

class Gram(nn.Module):
    def __init__(self):
        super(Gram, self).__init__()

    def forward(self, input):
        a, b, c, d = input.size()
        # 将特征图变换为 2 维向量
        feature = input.view(a * b, c * d)
        # 内积的计算方法其实就是特征图乘以它的逆
        gram = torch.mm(feature, feature.t())
        # 对得到的结果取平均值
        gram /= (a * b * c * d)
        return gram

gram = Gram()
gram

Gram 矩阵表示的是特征之间的相关程度,而这与图片的风格又有什么关系呢?关于 Gram 矩阵的更多数学理解可以参考 这篇回答。
接下来,让我们使用更加通俗的语言来阐述相关性和图片风格之间的关系。我们拿梵高的星空图举例:

image.png

这里我们要对梵高的星空图进行风格提取。假设,神经网络的某一层,有一个滤波器专门检测像“尖尖的塔顶”一样的东西。另一个滤波器专门检测黑色。又有一个滤波器负责检测圆圆的东西,又有一个滤波器用来检测金黄色。对梵高的原图做 Gram 矩阵,其中哪些特征的相关性矩阵会比较大呢?
如上图所示,“尖尖的”和“黑色”总是一起出现的,它们的相关性比较高。而“圆圆的”和“金黄色”也都是成对出现的,因此他们的相关性也比较高。而相对的, “尖尖的” 和 “金黄色” 的相关性就比较差。
因此在风格迁移的时候,计算机其实就是在内容图(待风格转换的图)里去寻找这种“匹配”,将尖尖的渲染为黑色(如塔尖),将圆圆的渲染为金黄色(如近圆的房顶)。
简单的说,由于图像的艺术风格就是其基本形状与色彩的组合方式,因此 Gram 矩阵能够很好的表征图像的艺术风格。
至此,我们可以对任意一张图片或者神经网络中的任一层的输出图层进行风格提取进而得到该图层的风格。

现在让我们来提取我们的风格图片 S 的风格:

target = gram(style_img)
# 此时 style_img 的通道为3 所以产生的风格特征为 3×3
target

在得到风格后,我们就需要计算风格损失了。这里我们可以采用交叉熵损失来计算任意两张图层的风格损失。风格损失函数的代码如下:

class Style_Loss(nn.Module):
    def __init__(self, target, weight):
        super(Style_Loss, self).__init__()
        # weight 和内容函数相似,表示的是权重 beta
        self.weight = weight
        # targer 表示图层目标。即新图像想要拥有的风格
        # 即保存目标风格
        self.target = target.detach() * self.weight
        self.gram = Gram()
        self.criterion = nn.MSELoss()

    def forward(self, input):
        # 加权计算 input 的 Gram 矩阵
        G = self.gram(input) * self.weight
        # 计算真实的风格和想要得到的风格之间的风格损失
        self.loss = self.criterion(G, self.target)
        out = input.clone()
        return out
    # 向后传播

    def backward(self, retain_variabels=True):
        self.loss.backward(retain_graph=retain_variabels)
        return self.loss

同样,让我们传入风格目标 target,得到一个可以计算任何图像与 target 之间的风格差异的损失函数:

# 传入模型所需参数
sl = Style_Loss(target, 1000)
# 传入一张随机图片进行测试
rand_img = torch.randn(style_img.data.size(), device=device)
# 损失函数层向前传播,进而得到损失
sl.forward(rand_img)
sl.loss

由于这里对风格损失赋予的权重是 1000 ,所以得到的风格损失较大。
至此,我们已经得到了风格损失和内容损失的具体表现形式。接下来,我们将会建立合适的神经网络模型并且阐述风格迁移的具体实现算法。

神经网络模型

VGG19
迁移算法主要依靠的网络结构是 VGG19 网络,这种网络结构和 VGG16 类似,也是神经网络中使用较为广泛的网络结构之一。由于本篇课程着重讲述的是风格迁移算法的实现。又因为 PyTorch 的官方工具包中已经为我们定义好了这种结构,因此这里就不对 VGG19 的网络结构做更深层的阐述了。如果想了解更多可以 参看这篇文章 。

让人感到高兴的是,PyTorch 官方工具包不仅为我们提供了 VGG19 的网络结构接口,还为我们提供了相应的预训练模型。我们可以通过 models.vgg19(pretrained=True).features 获得 VGG19 的所有池化层和卷积层的结构以及权重值。当 pretrained=True 时,计算机会先在本地的默认文件夹中寻找预训练模型并加载。如果没有找到,就会从身在遥远的大洋彼岸的服务器上下载该预训练模型。因此,在第一次加载 VGG19 时会非常非常非常慢(因为从外网下载东西的速度很慢)。
为了解决这一问题,我们将预训练模型上传到了课程的云存储中,并通过对 torch.utils.model_zoo.load_url 参数的设置,来让 Pytorch 不去遥远的彼岸下载预训练模型,而是通过我们指定的网址下载:

import torchvision.models as models
# 设置与预训练模型所在连接
torch.utils.model_zoo.load_url("https://labfile.oss.aliyuncs.com/courses/861/vgg19_pre.zip")

接下来,让我们来加载 VGG19 的网络结构和所对应的权重:

vgg = models.vgg19(pretrained=True).features
vgg = vgg.to(device)
vgg

风格迁移的网络模型
接下来,就让我们使用 VGG19 的网络结构、风格损失函数以及内容损失函数来构造一个用于图像风格迁移的神经网络模型。该网络模型只用到了 VGG19 的前 5 个卷积层。当然,你也可以根据自己的实际情况对模型结构进行优化,结构如下:

image.png

如上图所示,我们对每个网络层的输出都进行了一次风格提取,然后计算风格损失 Style_Loss(缩写为 SL),一共计算了 5 次风格损失。然后,利用 conv 4 的输出计算了一次内容损失(缩写为 CL)。而我们训练的最终目标就是让这些损失的加权和(即总损失)最小。
接下来让我们用代码来构造上面的网络结构:

content_layers_default = ['conv_4']
style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']
# 初始化一个 空的神经网络 model
model = nn.Sequential()
model = model.to(device)
# 构造网络模型,并且返回这些损失函数


def get_style_model_and_loss(style_img, content_img, cnn=vgg, style_weight=1000, content_weight=1,
                             content_layers=content_layers_default,
                             style_layers=style_layers_default):
    # 用列表来存上面6个损失函数
    content_loss_list = []
    style_loss_list = []

    # 风格提取函数
    gram = Gram()
    gram = gram.to(device)

    i = 1
    # 遍历 VGG19 ,找到其中我们需要的卷积层
    for layer in cnn:
        # 如果 layer 是  nn.Conv2d 对象,则返回 True
        # 否则返回 False
        if isinstance(layer, nn.Conv2d):
            # 将该卷积层加入我们的模型中
            name = 'conv_' + str(i)
            model.add_module(name, layer)

            # 判断该卷积层是否用于计算内容损失
            if name in content_layers_default:
                # 这里是把目标放入模型中,得到该层的目标
                target = model(content_img)
                # 目标作为参数传入具体的损失类中,得到一个工具函数。
                # 该函数可以计算任何图片与目标的内容损失
                content_loss = Content_Loss(target, content_weight)
                model.add_module('content_loss_' + str(i), content_loss)
                content_loss_list.append(content_loss)

            # 和内容损失相似,不过增加了一步:提取风格
            if name in style_layers_default:
                target = model(style_img)
                target = gram(target)
                # 目标作为参数传入具体的损失类中,得到一个工具函数。
                # 该函数可以计算任何图片与目标的风格损失
                style_loss = Style_Loss(target, style_weight)
                model.add_module('style_loss_' + str(i), style_loss)
                style_loss_list.append(style_loss)

            i += 1
        # 对于池化层和 Relu 层我们直接添加即可
        if isinstance(layer, nn.MaxPool2d):
            name = 'pool_' + str(i)
            model.add_module(name, layer)

        if isinstance(layer, nn.ReLU):
            name = 'relu' + str(i)
            model.add_module(name, layer)
    # 综上:我们得到了:
    # 一个具体的神经网络模型,
    # 一个风格损失函数集合(其中包含了 5 个不同风格目标的损失函数)
    # 一个内容损失函数集合(这里只有一个,你也可以多定义几个)
    return model, style_loss_list, content_loss_list

上面的代码构造了我们需要的网络模型并且对不同损失函数进行了归类。
同一类的损失函数可能也存在多组目标(这些目标是同一张图片在神经网络的不同层的表现形式)。由于不同网络层的感受野不同,它们识别到的特征也不同。因此我们对每一层的风格都进行了提取,进而得到了 5 个目标风格不同的损失函数(可以理解为 5 个不同的人对同一张图片的风格的见解)。当然内容损失也可以多定义几个,这里我们就不做尝试了。
我们可以传入风格图片和内容图片进行测试:

model, style_loss_list, content_loss_list = get_style_model_and_loss(
    style_img, content_img)
model

上面结果可以看出,前 5 个卷积层受到了改造,且后面的模型训练也只会用到前 5 个卷积层。
但是,为什么不直接把 VGG19 的其它部分(即后面部分)删除掉呢?这样做的目的是为了让我们可以更加简单的,通过修改 content_layers 和 style_layers 参数来修改整个模型的结构,增加更多的风格损失和内容损失,进而还原图片的更多细节。
模型的训练
为了方便讲解,我们将上面的网络结构,再次展示出来:

image.png

虽然我们现在建立了一个网络模型,但是还有很多东西我们并没说清楚。最重要的就是,该模型的 输入 与 输出 。
为什么我要把上面两个概念加重颜色呢? 因为这正是该网络结构的巧妙之处,也是该网络结构与其他我们做过的深度学习任务的不同之处。
该网络结构的输入其实是有三个,即风格图像 S 、内容图像 C 和随机图像 G 。其中,随机图像可以理解为我们使用随机数产生的一个没有任何意义的噪点图像。
传统深度学习任务的网络结构输出即为所求,但是本任务不同。该网络结构的输出其实就是一些特征,这些特征仅仅是用来计算损失的,无法称之为真正的图像(上图可知,有 5+1 个输出)。那么我们是怎样获得最后的,具有指定风格和指定内容的新图像的呢?
其实上图可以很好的阐述该模型一次正向传播的过程了,我们将 S,G,C 都放入上面的模型中进行计算。通过 S 和 G 计算出风格总损失,通过 G 和 C 计算内容总损失,进而得到模型的总损失。然后再利用总损失进行反向传播,并且利用梯度下降算法 调节 G 中的值 。
你没有看错,我们将模型进行训练,调节的不是神经网络层的权重(换句话说,网络结构的权重至始至终没有发生变化),调节的是 G 中的值。
我认为这才是整个图像风格迁移中最有意思且最难懂的地方。然而不幸的是,很多网上教程都忽略了这一点,让人以为整个网络模型的输出就是新图像。而事实是,模型所调节的参数才是新图像(参数即图像)。通过神经网络的后向传播,直接対新图像 G 的每个像素点的值进行调节,进而得到最符合期望的新图像 G。
换句话说,对于这种传统的图像风格迁移算法,保存模型是没有意义的,神经网络层中的参数在训练前和训练后并未发生变化。因为整个模型训练中,我们调节的都是 G ,而非模型中的参数。
因此,在定义优化器时,我们不能像传统深度学习一样,传入 model.params 。这里我们传入的应该是 G ,代码如下:

import torch.optim as optim


def get_input_param_optimier(input_img):
    # 将input_img的值转为神经网络中的参数类型
    input_param = nn.Parameter(input_img.data)
    # 告诉优化器,我们优化的是 input_img 而不是网络层的权重
    # 采用 LBFGS 优化器
    optimizer = optim.LBFGS([input_param])
    return input_param, optimizer


# 输入一个随机图片进行测试
get_input_param_optimier(rand_img)

接下来,就是模型的训练函数了。让我们按照上面的思路,对模型的训练函数进行编写:

# 传入的 input_img 是 G 中每个像素点的值,可以为一个随机图片
def run_style_transfer(content_img, style_img, input_img, num_epoches):
    print('Building the style transfer model..')
    # 指定所需要优化的参数,这里 input_param就是G中的每个像素点的值
    input_param, optimizer = get_input_param_optimier(input_img)

    print('Opimizing...')
    epoch = [0]
    while epoch[0] < num_epoches:
        # 这里我们自定义了总损失的计算方法
        def closure():
            input_param.data.clamp_(0, 1)  # 更新图像的数据
            # 将此时的 G 传入模型中,得到每一个网络层的输出
            model(input_param)
            style_score = 0
            content_score = 0
            # 清空之前的梯度
            optimizer.zero_grad()
            # 计算总损失,并得到各个损失的梯度
            for sl in style_loss_list:
                style_score += sl.backward()
            for cl in content_loss_list:
                content_score += cl.backward()

            epoch[0] += 1
            # 这里每迭代一次就进行一次输出
            # 你可以根据自身情况进行调节
            if epoch[0] % 1 == 0:
                print('run {}/80'.format(epoch))
                print('Style Loss: {:.4f} Content Loss: {:.4f}'.format(
                    style_score.data.item(), content_score.data.item()))
                print()

            return style_score + content_score
        # 更新 G
        optimizer.step(closure)
    # 返回训练完成的 G,此时的 G
    return input_param.data

最后,让我们调用上面的函数,正式开始进行模型的训练。由于内容图片 C 和风格图片 S 已经在上面定义好了。因此,我们只需要以随机噪点的方式初始化图像 G 即可。当然,除了直接以随机化的方式初始化图像 G,还有一种较为好的方式初始化 G。
这种方式就是将内容图像 C 中的每个像素点的值全部复制给图像 G。也就是说,新图像 G 中的每个像素点的初始值和 C 一致。这样有一个好处,就是可以减少模型的迭代次数。也就是说,模型训练开始时,图片 G 和内容图片 C 完全一致,我们只需要在尽量保留 G 的原内容的情况下,修改 G 的风格即可。代码如下(由于 CPU 运行,下面代码会运行 25 min 左右,请耐心等待):

# 初始化 G
input_img = content_img.clone()
# 进行模型训练,并且返回图片
out = run_style_transfer(content_img, style_img, input_img, num_epoches=80)
# 将图片转换成可 PIL 类型,便于展示
new_pic = transforms.ToPILImage()(out.cpu().squeeze(0))
print("训练完成")

训练完成后,让我们来加载一下这张图片:

# 展示图片
plt.imshow(new_pic)

我们对上面的模型进行了 80 次训练,结果如下,可以看到,新图像的效果还可以。

image.png

如果你想要获得更加好的结果,可以使用 GPU 进行更多的训练(一般为 80~200 之间,如果训练次数过多,可能会造成过拟合现象)。为了能够方便同学们更加快速的得到结果,这里我将代码放在了 Kaggle 上,你可以通过 该链接 访问相应的代码,并且利用 Kaggle 所提供的免费 GPU 进行运行。
提示:如果你想要在线上迁移自定义的图片,你可以先将需要迁移的图片上传到互联网中的任意一个图床中,然后使用 wget 将其下载到课程本地,重新运行上面代码即可(当然,需要修改上面代码中的图片路径)。

你可能感兴趣的:(25PyTorch 实现图像风格迁移)