PyTorch风格迁移，人人都是名画大师

爱好绘画的小伙伴们有没有想过将各种名画的风格融入自己的绘画作品当中？如今借助深度学习技术，很容易就能将名画的风格迁移到任何一张画中。

Neural Transfer

网络接收三张图片作为输入，一张内容图像，一张风格图像，一张由内容图像初始化的图像（最终将风格迁移到这张图像上来）。

损失函数

这里将会有两个损失函数：

1. 与内容图像之间的损失
2. 与风格图像之间的损失

内容损失

这个损失比较简单，只是单单衡量两张图片中的内容相差情况，因此使用MSE损失就好，一般损失中的target不能求梯度，需要detach：

class ContentLoss(nn.Module):

    def __init__(self, target,):
        super(ContentLoss, self).__init__()
        # we 'detach' the target content from the tree used
        # to dynamically compute the gradient: this is a stated value,
        # not a variable. Otherwise the forward method of the criterion
        # will throw an error.
        self.target = target.detach()

    def forward(self, input):
        self.loss = F.mse_loss(input, self.target)
        return input

注意，这其实并不是真正的PyTorch损失函数，如果要定义真正的内容损失函数，那么必须自己实现一个求导函数并且在backward中手动计算梯度。

风格损失

怎么来计算两张图片之间的风格差异的呢，我们可以先计算格拉姆矩阵（具体内容请参考Neural-Style algorithm的论文），格拉姆矩阵是一个矩阵与它的转置矩阵的积，相乘之后格拉姆矩阵必须进行归一化，每个元素除以所有元素的个数，这是因为相乘之后的格拉姆矩阵中可能存在很大的值，这些值会在梯度下降的过程中对前面的网络层有更大影响，而风格特征很可能出现在更深的网络层中，因此这种归一化操作是十分关键的。

def gram_matrix(input):
    a, b, c, d = input.size()  # a=batch size(=1)
    # b=number of feature maps
    # (c,d)=dimensions of a f. map (N=c*d)

    features = input.view(a * b, c * d)  # resise F_XL into \hat F_XL

    G = torch.mm(features, features.t())  # compute the gram product

    # we 'normalize' the values of the gram matrix
    # by dividing by the number of element in each feature maps.
    return G.div(a * b * c * d)

class StyleLoss(nn.Module):

    def __init__(self, target_feature):
        super(StyleLoss, self).__init__()
        self.target = gram_matrix(target_feature).detach()

    def forward(self, input):
        G = gram_matrix(input)
        self.loss = F.mse_loss(G, self.target)
        return input

导入模型

接下来使用VGG19提取特征，并且在多个卷积层之后添加计算内容损失和风格损失的模块：

        if name in content_layers:
            # add content loss:
            target = model(content_img).detach()
            content_loss = ContentLoss(target)
            model.add_module("content_loss_{}".format(i), content_loss)
            content_losses.append(content_loss)

        if name in style_layers:
            # add style loss:
            target_feature = model(style_img).detach()
            style_loss = StyleLoss(target_feature)
            model.add_module("style_loss_{}".format(i), style_loss)
            style_losses.append(style_loss)

与训练网络不同的是，这里并不是训练网络各个层的参数，而是对输入图像进行训练并最小化内容损失和风格损失，因此需要把输入图像设定为可求导：

optimizer = optim.LBFGS([input_img.requires_grad_()])

现在我们将毕加索的风格迁移到一张图片上看看。

内容图片：

风格图片：

风格迁移结果：

再试试把《神奈川冲浪里》的风格迁移到这张画上。

风格图片：

风格迁移结果：

结语

大家觉得人工智能学习图像风格的效果怎么样呢，各位小伙伴如果有自己喜欢的名画，不妨试试将它的风格迁移到自己的画作上来。

扫码关注微信公众号：机器工匠，回复关键字“风格迁移”获取完整代码、图片、论文。