神经风格转换(Style Transfer)小试牛刀

本文代码和部分内容参考课程:《动手学深度学习》:样式迁移
深度学习框架:MXNET(Python调用)
神经风格转换论文原文参考:A Neural Algorithm of Artistic Style

1 什么是神经风格转换

神经风格转换(也可称作“样式迁移”)是一种使用卷积神经网络自动将某一图像的样式(风格)应用到另一张图片上的技术,可以看做将某张图片自动施加滤镜的修饰技术。比如,我们可以把一张图片变为素描风格,油画风格,蜡笔风格等任何你想要的艺术风格。用神经网络实现这样一个小功能不得不说是一个挺有趣的事情。下面先展示我自己完成的一些例子:

神经风格转换(Style Transfer)小试牛刀_第1张图片 神经风格转换(Style Transfer)小试牛刀_第2张图片 神经风格转换(Style Transfer)小试牛刀_第3张图片
神经风格转换(Style Transfer)小试牛刀_第4张图片 神经风格转换(Style Transfer)小试牛刀_第5张图片 神经风格转换(Style Transfer)小试牛刀_第6张图片
神经风格转换(Style Transfer)小试牛刀_第7张图片 神经风格转换(Style Transfer)小试牛刀_第8张图片 神经风格转换(Style Transfer)小试牛刀_第9张图片
神经风格转换(Style Transfer)小试牛刀_第10张图片 神经风格转换(Style Transfer)小试牛刀_第11张图片 神经风格转换(Style Transfer)小试牛刀_第12张图片

上面系列图靠左的图为内容图片(Content),居中的为风格图片(Style),最右侧为输出,即我们期望获得的图像效果。从结果来看,生成的图片效果还不错。

2 神经风格转换实现原理

首先,我们初始化合成图像,例如将其初始化成内容图像。该合成图像是样式迁移过程中唯一需要更新的变量,即神经风格转换所需迭代的模型参数。
然后,我们选择一个预训练的卷积神经网络(论文原文和本文代码均使用了VGG-19)来抽取图像的特征,其中的模型参数在训练中无须更新。深度卷积神经网络凭借多个层逐级抽取图像的特征。我们可以选择其中某些层的输出作为内容特征或样式特征。以下图为例,这里选取的预训练的神经网络含有3个卷积层,其中第二层输出图像的内容特征,而第一层和第三层的输出被作为图像的样式特征。
接下来,我们通过正向传播(实线箭头方向)计算神经风格转的损失函数,并通过反向传播(虚线箭头方向)迭代模型参数,即不断更新合成图像。
样式迁移常用的损失函数通常由3部分组成:内容损失(content loss)使合成图像与内容图像在内容特征上接近,样式损失style loss)令合成图像与样式图像在样式特征上接近,而总变差损失total variation loss)使生产图像各个像素点与之周围像素点的差异值降低,有助于减少合成图像中的噪点。最后,当模型训练结束时,我们输出样式迁移的模型参数,即得到最终的合成图像。
神经风格转换(Style Transfer)小试牛刀_第13张图片

3 实现方法

3.1图片处理

首先我们导入必要的包和模块:

%matplotlib inline
import d2lzh as d2l
from mxnet import autograd, gluon, image, init, nd
from mxnet.gluon import model_zoo, nn
import time

其中d2lzh是《动手学深度学习》课程中提供的一个工具包,里面内置了大量深度学习中常用的函数,方便日常使用。
然后载入内容图片和风格图片:

content_img = image.imread('内容图片(content)的地址')
d2l.plt.imshow(content_img.asnumpy());

style_img = image.imread('风格图片(style)的地址')
d2l.plt.imshow(style_img.asnumpy());

定义图片预处理和后处理函数,预处理将图片RGB通道做标准化,并将结果变换成卷积神经网络接受的输入格式;后处理将神经网络的输出还原为图片原始格式:

# 图片RGB三通道的均值和标准差
'''这里我不是很懂,不同的图片这个值应该是不一样的,这个应该根据传入的图片分别
   计算各自的均值和标准差吧,原代码直接用给出的这个值标准化内容图片和风格图片,
   不过所有图片好像都可以用这个。。。而且效果还不错。。。在此提出疑问QAQ。
'''
rgb_mean = nd.array([0.485, 0.456, 0.406])
rgb_std = nd.array([0.229, 0.224, 0.225])

def preprocess(img, image_shape):
    img = image.imresize(img, *image_shape)
    img = (img.astype('float32') / 255 - rgb_mean) / rgb_std
    return img.transpose((2, 0, 1)).expand_dims(axis=0)

def postprocess(img):
    img = img[0].as_in_context(rgb_std.context)
    return (img.transpose((1, 2, 0)) * rgb_std + rgb_mean).clip(0, 1)# clip用来确保像素值在0到1之间,因为神经网络可能会迭代一些负数出来哦

3.2 特征提取

和论文原文一样,使用了预训练的VGG-19神经网络,该网络使用了ImageNet数据集来训练,有空试试其他像ResNet系列,Inception系列网络的效果。
首先导入模型:

pretraind_VGG = model_zoo.vision.vgg19(pretrained=True, root='这里是你的模型参数缓存路径,不设置的话为默认值')
pretraind_VGG.features #查看一下网络的构成(不考虑output层))

设置我们需要用到的卷积层,来提取我们的特征。这里我提取风格特征的层用了5个,提取内容特征的层只用了1个,且输出内容特征的层尽量靠后,这样可以尽量避免生成的图像过多的保留内容图的细节。顺带提一下,这里的层都用的是未激活的卷积层,有人使用网络中的relu激活层做特征提取,训练结果是否有差异值得研究

style_layers, content_layers = [0, 5, 10, 12, 19, 28], [29]

构造需要用到神经网络(把所能用到的最大层数后面的网络层丢掉):

net = nn.HybridSequential()
for i in range(max(content_layers + style_layers) + 1):
    net.add(pretraind_VGG.features[i])

由于神经网络的向前传播只能输出最后一层的值,这里我们需要获得每一层的输出,因此定义特征提取函数:

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

定义好特征提取函数后,我们可以用构造的网络一次性将内容特征和风格特征提取出来放入各自列表中(因为这些值相当于参数,不会拿来训练,所以也不会改变),定义两个函数获得特征层和内容层:

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

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

3.3 定义损失函数

定义内容损失函数(简单的最小二乘,不再解释):

def content_loss(Y_hat, Y):
    return (Y_hat - Y).square().mean()

定义gram矩阵函数和风格损失函数(关于为什么gram矩阵能表达图片“风格”笔者也不是专业人士,只有自己一点浅薄的理解,从协方差的角度考虑可能该矩阵能反映图片像素之间的相关性分布?gram矩阵的公式及详细的介绍可以参考一下这篇博客:Gram Matrices理解):

def gram(X):
    num_channels, n = X.shape[1], X.shape[2]*X.shape[3]
    X = X.reshape((num_channels, n))
    return nd.dot(X, X.T) / X.size

def style_loss(Y_hat, gram_Y):
    return (gram(Y_hat) - gram_Y).square().mean()

定义总变差损失函数:
m i n ∑ i , j ∣ x i , j − x i + 1 , j ∣ + ∣ x i , j − x i , j + 1 ∣ min\sum_{i,j} \left|x_{i,j} - x_{i+1,j}\right| + \left|x_{i,j} - x_{i,j+1}\right| mini,jxi,jxi+1,j+xi,jxi,j+1
从公式容易看出,该式使得点 x i , j x_{i,j} xi,j与其右侧和下侧相邻的像素点差值尽可能小,这是一种有效的降噪手段。需要注意的是,在论文原文中作者没有计算总变差损失,因此这是后来对原有成像算法的改进:

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

将三种不同类型的损失加权平均,构造损失函数(这里面content_weight, style_weight和tv_weight都是超参,是各损失函数的权值):

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 = nd.add_n(*styles_l) + nd.add_n(*contents_l) + tv_l
    return contents_l, styles_l, tv_l, l

3.4 创建和初始化合成图像

合成图像在本项目中是待训练的参数,相当于普通神经网络的权值 W e i g h t Weight Weight以及偏置 B i a s Bias Bias,因此需要将合成图像构造成深度学习框架可识别的参数模型,在此继承mxnet.nn.HybridBlock类,定义一个GeneratedImage类用于创建合成图像:

class GeneratedImage(nn.HybridBlock):
    def __init__(self, img_shape, **kwargs):
        super(GeneratedImage, self).__init__(**kwargs)
        self.weight = self.params.get('weight', shape=img_shape)

    def forward(self):
        return self.weight.data()

然后定义初始化函数,用于将创建的合成图像初始化。这里我们还同时把训练器定义了(我们直接将图像初始化成内容图像,有的例子还在其基础上加了随机白噪声,大家可以测试一下各自效果):

def get_inits(X, ctx, lr, styles_Y):
    gen_img = GeneratedImage(X.shape)
    gen_img.initialize(init.Constant(X), ctx=ctx, force_reinit=True)
    trainer = gluon.Trainer(gen_img.collect_params(), 'adam', {'learning_rate': lr})
    styles_Y_gram = [gram(Y) for Y in styles_Y]
    return gen_img(), styles_Y_gram, trainer

3.5 定义训练函数

这里不多介绍,MXNET标准的定义方法,只不过我们额外加了一个时间戳,并加入了梯度衰减,每50次我们输出一次模型各损失函数的值:

def train(X, contents_Y, styles_Y, ctx, lr, max_epochs, lr_decay_epoch):
    X, styles_Y_gram, trainer = get_inits(X, ctx, lr, styles_Y)
    for i in range(max_epochs):
        start = time.time()
        with autograd.record():
            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(1)
        nd.waitall()
        if i % 50 == 0 and i != 0:
            print('epoch %3d, content loss %.2f, style loss %.2f, '
                  'TV loss %.2f, %.2f sec'
                  % (i, nd.add_n(*contents_l).asscalar(),
                     nd.add_n(*styles_l).asscalar(), tv_l.asscalar(),
                     time.time() - start))
        if i % lr_decay_epoch == 0 and i != 0:
            trainer.set_learning_rate(trainer.learning_rate * 0.1)
            print('change lr to %.1e' % trainer.learning_rate)
    return X

3.6 训练并输出结果

训练这个合成图像时,笔者用了阶梯式训练的方法,即先将图片用较小尺寸快速训练出一个结果,再将所得图片增大,作为我们下一次训练的初始值进行训练,不断重复这个过程,直到获得我们目标尺寸的图像输出。
笔者在实验中发现,小尺寸的图片训练迭代速度非常快,但由于图片的压缩使得生成图片细节模糊;而对大尺寸的图片进行训练时比较耗时,因此先用小尺寸图片训练出的结果作为大尺寸图片的输入时,可一定程度上优化大尺寸图片训练的初始值,从而减少训练次数,从整体上减少耗时。
第一次将图片压缩为(150, 225)尺寸训练,迭代800次,学习率设置为0.05,梯度衰减区间设置为200:

# 尺寸自己设定
content_weight, style_weight, tv_weight, ctx, image_shape= 1, 3e3, 30, d2l.try_gpu(), (150, 225)
net.collect_params().reset_ctx(ctx)
net.hybridize()
content_X, contents_Y = get_contents(image_shape, ctx)
_, styles_Y = get_styles(image_shape, ctx)
output = train(content_X, contents_Y, styles_Y, ctx, 0.05, 800, 200)

第二次将输出结果增大到(400, 600),迭代400次,学习率设置为0.02,梯度衰减区间设置为200:

# 尺寸自己设定
image_shape = (400, 600)
_, content_Y = get_contents(image_shape, ctx)
_, style_Y = get_styles(image_shape, ctx)
X = preprocess(postprocess(output) * 255, image_shape)
output = train(X, content_Y, style_Y, ctx, 0.02, 400, 200)

最后将输出结果增大到目标尺寸(这里我用的内容图片原始尺寸),迭代200次,学习率0.01,梯度衰减区间100:

# 尺寸自己设定
image_shape = (540, 800)
_, content_Y = get_contents(image_shape, ctx)
_, style_Y = get_styles(image_shape, ctx)
X = preprocess(postprocess(output) * 255, image_shape)
output = train(X, content_Y, style_Y, ctx, 0.01, 200, 100)

笔者GTX 1060的gpu整体耗时大概2分多钟能得到质量不错的输出,前提是图片不能太大(什么4K高清就算了。。。QAQ),结果还可以接受。
最后查看一下输出的结果,然后保存生成的图片就好啦:

d2l.plt.imshow(postprocess(output).asnumpy())
d2l.plt.imsave('你保存图片的地址.XXX.png(jpg...等等)', postprocess(output).asnumpy())

3.7 后记

通过这个小练习熟悉了神经风格转换的基本原理和实现方式,同时对gram矩阵有了一些自己的认识。使用了预训练的VGG-19模型,熟悉了迁移学习(更准确的说这只是使用了训练好的网络,还谈不上迁移学习,但二者对利用该网络对数据特征提取的方式是一致的)。最后,笔者觉得每获得一张图片就要使用神经网络训练多次效率还是太低了,能否像一般深度学习应用一样,预先使用数据集训练一个鲁棒的网络,使得我们可以将内容图片按需要的风格通过一次前向传播就获得结果是一个不错的思考。

你可能感兴趣的:(机器学习和深度学习)