【PyTorch】神经风格迁移项目

神经风格迁移中，取一个内容图像和一个风格图像，综合内容图像的内容和风格图像的艺术风格生成新的图像。

目录

准备数据

处理数据

神经风格迁移模型

加载预训练模型

定义损失函数

定义优化器

运行模型

---

准备数据

创建data文件夹，放入一张内容图片（左），一张风格图片（右），分别命名为content和style

from PIL import Image
path2content= "./data/content.jpg"
path2style= "./data/style.jpg"
content_img = Image.open(path2content)
style_img = Image.open(path2style)

处理数据

调用torchvision.transforms包中Resize、ToTensor和Normalize对图像进行预处理

import torchvision.transforms as transforms

h, w = 256, 384 
mean_rgb = (0.485, 0.456, 0.406)
std_rgb = (0.229, 0.224, 0.225)
transformer = transforms.Compose([
                    # 将图像缩放到指定大小
                    transforms.Resize((h,w)),  
                    # 将图像转换为张量
                    transforms.ToTensor(),
                    # 对图像进行标准化处理
                    transforms.Normalize(mean_rgb, std_rgb)])  

content_tensor = transformer(content_img)
print(content_tensor.shape, content_tensor.requires_grad)

style_tensor = transformer(style_img)
print(style_tensor.shape, style_tensor.requires_grad)

# 克隆content_tensor作为输入图像，并设置requires_grad为True，表示需要计算梯度
input_tensor = content_tensor.clone().requires_grad_(True)
print(input_tensor.shape, input_tensor.requires_grad)

import torch
from torchvision.transforms.functional import to_pil_image
# 将图像张量转换为所需PIL图像
def imgtensor2pil(img_tensor):
    # 克隆并分离图像张量
    img_tensor_c = img_tensor.clone().detach()
    # 将图像张量乘以标准RGB值
    img_tensor_c*=torch.tensor(std_rgb).view(3,1,1)
    # 将图像张量加上均值RGB值
    img_tensor_c+=torch.tensor(mean_rgb).view(3,1,1)
    # 将图像张量限制在0到1之间
    img_tensor_c = img_tensor_c.clamp(0,1)
    # 将图像张量转换为PIL图像
    img_pil=to_pil_image(img_tensor_c)
    # 返回PIL图像
    return img_pil

import matplotlib.pylab as plt
%matplotlib inline

plt.imshow(imgtensor2pil(content_tensor))
plt.title("content image");
plt.imshow(imgtensor2pil(style_tensor))
plt.title("style image");

神经风格迁移模型

保持模型参数不变，更新模型的输入

加载预训练模型

import torchvision.models as models
# 检查是否有可用的GPU，如果没有则使用CPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 加载预训练的VGG19模型，并将其特征提取部分移动到指定的设备上，并将其设置为评估模式
model_vgg = models.vgg19(pretrained=True).features.to(device).eval()
# 将模型的所有参数设置为不需要梯度，即不进行反向传播
for param in model_vgg.parameters():
    param.requires_grad_(False)   
print(model_vgg)

定义损失函数

# 定义函数，获取模型中指定层的特征
def get_features(x, model, layers):
    # 创建一个空字典，用于存储特征
    features = {}
    # 遍历模型的所有子层
    for name, layer in enumerate(model.children()):
        # 将输入数据传入子层，得到输出数据
        x = layer(x)
        # 如果子层的名称在指定的层列表中
        if str(name) in layers:
            # 将输出数据存储到字典中，键为子层的名称
            features[layers[str(name)]] = x
    # 返回字典
    return features

# 定义函数，于计算gram矩阵
def gram_matrix(x):
    # 获取输入张量的维度
    n, c, h, w = x.size()
    # 将输入张量展平
    x = x.view(n*c, h * w)
    # 计算gram矩阵
    gram = torch.mm(x, x.t())
    return gram

import torch.nn.functional as F

# 定义函数，获取内容损失
def get_content_loss(pred_features, target_features, layer):
    # 获取目标特征
    target= target_features[layer]
    # 获取预测特征
    pred = pred_features [layer]
    # 计算均方误差损失
    loss = F.mse_loss(pred, target)
    return loss

# 定义函数，获取风格损失
def get_style_loss(pred_features, target_features, style_layers_dict):  
    # 初始化损失为0
    loss = 0
    # 遍历style_layers_dict中的每一层
    for layer in style_layers_dict:
        # 获取预测特征
        pred_fea = pred_features[layer]
        # 计算预测特征的gram矩阵
        pred_gram = gram_matrix(pred_fea)
        # 获取预测特征的shape
        n, c, h, w = pred_fea.shape
        # 获取目标特征的gram矩阵
        target_gram = gram_matrix (target_features[layer])
        # 计算当前层的损失
        layer_loss = style_layers_dict[layer] *  F.mse_loss(pred_gram, target_gram)
        # 将当前层的损失加到总损失中
        loss += layer_loss/ (n* c * h * w)
    # 返回总损失
    return loss

# 定义特征层字典，用于存储不同层的特征
feature_layers = {'0': 'conv1_1',
                  '5': 'conv2_1',
                  '10': 'conv3_1',
                  '19': 'conv4_1',
                  '21': 'conv4_2',  
                  '28': 'conv5_1'}

# 将内容张量增加一个维度，并将其移动到指定设备上
con_tensor = content_tensor.unsqueeze(0).to(device)

sty_tensor = style_tensor.unsqueeze(0).to(device)

# 获取内容张量的特征
content_features = get_features(con_tensor, model_vgg, feature_layers)

style_features = get_features(sty_tensor, model_vgg, feature_layers)

# 遍历content_features字典中的所有key
for key in content_features.keys():
    # 打印每个key对应的值的形状
    print(content_features[key].shape)

定义优化器

from torch import optim

# 克隆con_tensor，并设置requires_grad_为True，表示需要计算梯度
input_tensor = con_tensor.clone().requires_grad_(True)
# 使用Adam优化器，优化input_tensor，学习率为0.01
optimizer = optim.Adam([input_tensor], lr=0.01)

运行模型

# 定义训练的轮数
num_epochs = 300
# 定义内容损失的权重
content_weight = 1e1
# 定义风格损失的权重
style_weight = 1e4
# 定义内容层
content_layer = "conv5_1"
# 定义风格层及其权重
style_layers_dict = { 'conv1_1': 0.75,
                      'conv2_1': 0.5,
                      'conv3_1': 0.25,
                      'conv4_1': 0.25,
                      'conv5_1': 0.25}

# 遍历每一轮
for epoch in range(num_epochs+1):
    # 梯度清零
    optimizer.zero_grad()
    # 获取输入特征
    input_features = get_features(input_tensor, model_vgg, feature_layers)
    # 获取内容损失
    content_loss = get_content_loss (input_features, content_features, content_layer)
    # 获取风格损失
    style_loss = get_style_loss(input_features, style_features, style_layers_dict)
    # 计算神经损失
    neural_loss = content_weight * content_loss + style_weight * style_loss
    # 反向传播
    neural_loss.backward(retain_graph=True)
    # 更新参数
    optimizer.step()
    
    # 每隔100轮打印一次损失
    if epoch % 100 == 0:
        print('epoch {}, content loss: {:.2}, style loss {:.2}'.format(
          epoch,content_loss, style_loss))

打印输出图片（左），对比原始内容图片（右）

plt.imshow(imgtensor2pil(input_tensor[0].cpu()));