VGG网络实现图片的风格迁移

from torchvision import models from torchvision import transforms from PIL import Image import torch import torch.nn as nn import numpy as np import matplotlib.pyplot as plt import imageio import cv2 #调用GPU device = torch.device("cuda" if torch.cuda.is_available() else "cpu") #图片初始化 def load_image(image_path, transform=None, max_size=None, shape=None): image = Image.open(image_path) if max_size: scale = max_size / max(image.size) # 限制大小 size = np.array(image.size) * scale image = image.resize(size.astype(int), Image.ANTIALIAS) if shape: image = image.resize(shape, Image.LANCZOS) if transform: image = transform(image).unsqueeze(0) # [1,3,h,w] return image.to(device) transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) content = load_image("./content.jpeg", transform, max_size=400) style = load_image("./style.jpg", transform, shape=[content.size(3), content.size(2)]) unloader = transforms.ToPILImage() def imshow(tensor): image = tensor.cpu().clone() image = image.squeeze(0) image = unloader(image) image.show() #定义VGG网络并将数值放入GPU进行计算 class VGGNet(nn.Module): def __init__(self): super(VGGNet, self).__init__() self.select = ["0", "5", "10", "19", "28"] self.vgg = models.vgg19(pretrained=True).features def forward(self, x): features = [] for name, layer in self.vgg._modules.items(): x = layer(x) if name in self.select: features.append(x) return features vgg = VGGNet().to(device).eval() #开始训练进行风格迁移,每100轮展现一次 target = content.clone().requires_grad_(True) optimizer = torch.optim.Adam([target],lr = 0.003,betas=[0.5,0.999]) #开启多次迭代，每100次显示一次 total_steps = 3000 for step in range(total_steps): target_features = vgg(target) content_features = vgg(content) style_features = vgg(style) content_loss = style_loss = 0. for f1, f2, f3 in zip(target_features, content_features, style_features): content_loss += torch.mean((f1 - f2) ** 2) _, c, h, w = f1.size() f1 = f1.view(c, h * w) # c,h*w f3 = f3.view(c, h * w) f1 = torch.mm(f1, f1.t()) # 点积运算，需要转置 c*c f3 = torch.mm(f3, f3.t()) style_loss += torch.mean((f1 - f3) ** 2) / (c * h * w) loss = content_loss + style_loss * 100. optimizer.zero_grad() loss.backward() optimizer.step() if step % 100 == 0: print("step{}/{},content loss:{:4f},style loss:{:4f}". format(step, total_steps, content_loss.item(), style_loss.item())) denorm = transforms.Normalize([-2.12, -2.04, -1.80], [4.37, 4.46, 4.44]) img = target.clone().squeeze() img = denorm(img).clamp_(0, 1) img = img.clamp_(0, 1) imshow(img)