PyTorch 进行神经传递

神经传递是一种强大的技术，它允许我们将一张图像的内容与另一张图像的风格结合起来，生成具有艺术效果的新图像。

一、基本原理

神经传递算法的核心在于定义两个距离：内容距离和风格距离。

• 内容距离 ：衡量两张图像内容的差异。我们通过比较特征图的均方误差来计算内容距离。
• 风格距离 ：衡量两张图像风格的差异。我们使用 Gram 矩阵来捕捉图像的风格特征，然后计算 Gram 矩阵之间的均方误差作为风格距离。

算法的目标是找到一张输入图像，使其与内容图像的内容距离和与风格图像的风格距离都尽可能小。

二、导入必要的软件包并选择设备

我们首先导入实现神经传递所需的软件包，并选择运行设备。

from __future__ import print_function


import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim


from PIL import Image
import matplotlib.pyplot as plt


import torchvision.transforms as transforms
import torchvision.models as models


import copy


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

三、加载图像

我们需要加载内容图像和风格图像，并将它们转换为 PyTorch 张量。同时，我们还需要一个函数来显示图像。

imsize = 512 if torch.cuda.is_available() else 128  # 如果没有 GPU，使用较小的图像尺寸


loader = transforms.Compose([
    transforms.Resize(imsize),
    transforms.ToTensor()])


def image_loader(image_name):
    image = Image.open(image_name)
    image = loader(image).unsqueeze(0)  # 添加批量维度
    return image.to(device, torch.float)


style_img = image_loader("style_image.jpg")  # 风格图像路径
content_img = image_loader("content_image.jpg")  # 内容图像路径


assert style_img.size() == content_img.size(), "内容图像和风格图像的尺寸必须相同"


unloader = transforms.ToPILImage()  # 将张量转换为 PIL 图像


def imshow(tensor, title=None):
    image = tensor.cpu().clone()
    image = image.squeeze(0)
    image = unloader(image)
    plt.imshow(image)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)


plt.figure()
imshow(style_img, title='Style Image')


plt.figure()
imshow(content_img, title='Content Image')

四、定义损失函数

我们定义内容损失和风格损失函数。

class ContentLoss(nn.Module):
    def __init__(self, target):
        super(ContentLoss, self).__init__()
        self.target = target.detach()


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


def gram_matrix(input):
    a, b, c, d = input.size()
    features = input.view(a * b, c * d)
    G = torch.mm(features, features.t())
    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 模型，并在适当的位置插入内容损失和风格损失层。

cnn = models.vgg19(pretrained=True).features.to(device).eval()


cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device)
cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device)


class Normalization(nn.Module):
    def __init__(self, mean, std):
        super(Normalization, self).__init__()
        self.mean = mean.view(-1, 1, 1)
        self.std = std.view(-1, 1, 1)


    def forward(self, img):
        return (img - self.mean) / self.std


content_layers_default = ['conv_4']
style_layers_default = ['conv_1', 'conv_2', 'conv_3', 'conv_4', 'conv_5']


def get_style_model_and_losses(cnn, normalization_mean, normalization_std, style_img, content_img):
    cnn = copy.deepcopy(cnn)
    normalization = Normalization(normalization_mean, normalization_std).to(device)
    model = nn.Sequential(normalization)
    content_losses = []
    style_losses = []
    i = 0
    for layer in cnn.children():
        if isinstance(layer, nn.Conv2d):
            i += 1
            name = 'conv_{}'.format(i)
        elif isinstance(layer, nn.ReLU):
            name = 'relu_{}'.format(i)
            layer = nn.ReLU(inplace=False)
        elif isinstance(layer, nn.MaxPool2d):
            name = 'pool_{}'.format(i)
        elif isinstance(layer, nn.BatchNorm2d):
            name = 'bn_{}'.format(i)
        else:
            raise RuntimeError('未识别的层：{}'.format(layer.__class__.__name__))


        model.add_module(name, layer)


        if name in content_layers_default:
            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_default:
            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)


    for i in range(len(model) - 1, -1, -1):
        if isinstance(model[i], ContentLoss) or isinstance(model[i], StyleLoss):
            break


    model = model[:(i + 1)]
    return model, style_losses, content_losses

六、选择输入图像并定义优化器

我们可以选择使用内容图像或白噪声作为输入图像，并定义优化器。

input_img = content_img.clone().requires_grad_(True)
## input_img = torch.randn(content_img.data.size(), device=device).requires_grad_(True)


def get_input_optimizer(input_img):
    optimizer = optim.LBFGS([input_img])
    return optimizer


optimizer = get_input_optimizer(input_img)

七、运行神经传递

最后，我们定义并运行神经传递函数，生成具有指定艺术风格的内容图像。

def run_style_transfer(cnn, normalization_mean, normalization_std, content_img, style_img, input_img, num_steps=300, style_weight=1000000, content_weight=1):
    print('构建神经传递模型...')
    model, style_losses, content_losses = get_style_model_and_losses(cnn, normalization_mean, normalization_std, style_img, content_img)
    optimizer = get_input_optimizer(input_img)
    print('开始优化...')
    run = [0]
    while run[0] <= num_steps:
        def closure():
            input_img.data.clamp_(0, 1)
            optimizer.zero_grad()
            model(input_img)
            style_score = 0
            content_score = 0
            for sl in style_losses:
                style_score += sl.loss
            for cl in content_losses:
                content_score += cl.loss
            style_score *= style_weight
            content_score *= content_weight
            loss = style_score + content_score
            loss.backward()
            run[0] += 1
            if run[0] % 50 == 0:
                print("运行 {}:".format(run))
                print('风格损失 : {:4f} 内容损失: {:4f}'.format(style_score.item(), content_score.item()))
                print()
            return style_score + content_score
        optimizer.step(closure)
    input_img.data.clamp_(0, 1)
    return input_img


output = run_style_transfer(cnn, cnn_normalization_mean, cnn_normalization_std, content_img, style_img, input_img)
plt.figure()
imshow(output, title='Output Image')
plt.ioff()
plt.show()

通过本教程，你成功地在 PyTorch 中实现了神经传递，将内容图像和风格图像结合在一起，生成了具有艺术风格的新图像。在编程狮（W3Cschool）网站上，你可以找到更多关于 PyTorch 的详细教程和实战案例，帮助你进一步提升深度学习技能，成为人工智能领域的编程大神。