神经风格迁移中,取一个内容图像和一个风格图像,综合内容图像的内容和风格图像的艺术风格生成新的图像。
目录
准备数据
处理数据
神经风格迁移模型
加载预训练模型
定义损失函数
定义优化器
运行模型
准备数据
创建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 transformsh, 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_pilimport matplotlib.pylab as plt
%matplotlib inlineplt.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 gramimport 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): # 初始化损失为0loss = 0# 遍历style_layers_dict中的每一层for layer in style_layers_dict:# 获取预测特征pred_fea = pred_features[layer]# 计算预测特征的gram矩阵pred_gram = gram_matrix(pred_fea)# 获取预测特征的shapen, 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()));