第一种代码
from __future__ import division
from torchvision import models
from torchvision import transforms
from PIL import Image
import argparse
import torch
import torchvision
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")#选择用gpu还是cpu
print(device)
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)
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])
]) # 来自ImageNet的mean和variance
content = load_image("D:/电脑壁纸/v2-8c48ebffbc108fb45f04e5ff0bff4015_r.jpg", transform, max_size=400)
style = load_image("D:/转移事项/桌面/R.jpg", transform, shape=[content.size(2), content.size(3)])
unloader = transforms.ToPILImage() # reconvert into PIL image
def imshow(tensor, title=None):
image = tensor.cpu().clone() # we clone the tensor to not do changes on it
image = image.squeeze(0) # remove the fake batch dimension
image = unloader(image)
plt.imshow(image)
if title is not None:
plt.title(title)
plt.pause(0.001) # pause a bit so that plots are updated
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
target = content.clone().requires_grad_(True)
optimizer = torch.optim.Adam([target], lr=0.003, betas=[0.5, 0.999])
vgg = VGGNet().to(device).eval()
target_features = vgg(target)
total_step = 2000
style_weight = 100.
for step in range(total_step):
target_features = vgg(target)
content_features = vgg(content)
style_features = vgg(style)
style_loss = 0
content_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)
f3 = f3.view(c, h * w)
# 计算gram matrix
f1 = torch.mm(f1, f1.t())
f3 = torch.mm(f3, f3.t())
style_loss += torch.mean((f1 - f3) ** 2) / (c * h * w)
loss = content_loss + style_weight * style_loss
# 更新target
optimizer.zero_grad()
loss.backward()
optimizer.step()
if step % 10 == 0:
print("Step [{}/{}], Content Loss: {:.4f}, Style Loss: {:.4f}"
.format(step, total_step, 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).clam99p_(0, 1)
imshow(img)
plt.show()
第二种代码
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from PIL import Image
from copy import deepcopy
import matplotlib.pyplot as plt
import numpy as np
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
class ContentLoss(nn.Module):
def __init__(self, target):
super(ContentLoss, self).__init__()
self.target = target.detach()
def forward(self, input):
self.loss = torch.sum((input-self.target) ** 2) / 2.0
return input
class StyleLoss(nn.Module):
def __init__(self, target_feature):
super(StyleLoss, self).__init__()
self.target = gram_matrix(target_feature).detach()
def forward(self, input):
a, b, c, d = input.size()
G = gram_matrix(input)
self.loss = torch.sum((G-self.target) ** 2) / (4.0 * b * b * c * d)
return input
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)
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])
]) # 来自ImageNet的mean和variance
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
content = load_image("D:/电脑壁纸/v2-8c48ebffbc108fb45f04e5ff0bff4015_r.jpg", transform, max_size=400)
style = load_image("D:/转移事项/桌面/R.jpg", transform, shape=[content.size(2), content.size(3)])
content_layers = ['conv_4_2'] # 内容损失函数使用的卷积层
style_layers = ['conv_1_1', 'conv2_1', 'conv_3_1', 'conv_4_1', 'conv5_1'] # 风格损失函数使用的卷积层
content_weights = [1] # 内容损失函数的权重
style_weights = [1e3, 1e3, 1e3, 1e3, 1e3] # 风格损失函数的权重
num_steps=1000 # 最优化的步数
class Model:
def __init__(self, device):
cnn = torchvision.models.vgg19(pretrained=True).features.to(device).eval()
self.cnn = deepcopy(cnn) # 获取预训练的VGG19卷积神经网络
self.device = device
self.content_losses = []
self.style_losses = []
def run(self):
content_image = load_image("D:/电脑壁纸/v2-8c48ebffbc108fb45f04e5ff0bff4015_r.jpg", transform, max_size=400)
style_image = load_image("D:/转移事项/桌面/R.jpg", transform, shape=[content.size(2), content.size(3)])
self._build(content_image, style_image) # 建立损失函数
output_image = self._transfer(content_image) # 进行最优化
return output_image
def _build(self, content_image, style_image):
self.model = nn.Sequential()
block_idx = 1
conv_idx = 1
# 逐层遍历VGG19，取用需要的卷积层
for layer in self.cnn.children():
# 识别该层类型并进行编号命名
if isinstance(layer, nn.Conv2d):
name = 'conv_{}_{}'.format(block_idx, conv_idx)
conv_idx += 1
elif isinstance(layer, nn.ReLU):
name = 'relu_{}_{}'.format(block_idx, conv_idx)
layer = nn.ReLU(inplace=False)
elif isinstance(layer, nn.MaxPool2d):
name = 'pool_{}'.format(block_idx)
block_idx += 1
conv_idx = 1
elif isinstance(layer, nn.BatchNorm2d):
name = 'bn_{}'.format(block_idx)
else:
raise Exception("invalid layer")
self.model.add_module(name, layer)
# 添加内容损失函数
if name in content_layers:
target = self.model(content_image).detach()
content_loss = ContentLoss(target)
self.model.add_module("content_loss_{}_{}".format(block_idx, conv_idx), content_loss)
self.content_losses.append(content_loss)
# 添加风格损失函数
if name in style_layers:
target_feature = self.model(style_image).detach()
style_loss = StyleLoss(target_feature)
self.model.add_module("style_loss_{}_{}".format(block_idx, conv_idx), style_loss)
self.style_losses.append(style_loss)
# 留下有用的部分
i = 0
for i in range(len(self.model) - 1, -1, -1):
if isinstance(self.model[i], ContentLoss) or isinstance(self.model[i], StyleLoss):
break
self.model = self.model[:(i + 1)]
def _transfer(self, content_image):
output_image = content_image.clone()
random_image = torch.randn(content_image.data.size(), device=self.device)
output_image = 0.4 * output_image + 0.6 * random_image
optimizer = torch.optim.LBFGS([output_image.requires_grad_()])
print('Optimizing..')
run = [0]
while run[0] <= num_steps:
def closure():
optimizer.zero_grad()
self.model(output_image)
style_score = 0
content_score = 0
for sl, sw in zip(self.style_losses, style_weights):
style_score += sl.loss * sw
for cl, cw in zip(self.content_losses, content_weights):
content_score += cl.loss * cw
loss = style_score + content_score
loss.backward()
run[0] += 1
if run[0] % 50 == 0:
print("iteration {}: Loss: {:4f} Style Loss: {:4f} Content Loss: {:4f}"
.format(run, loss.item(), style_score.item(), content_score.item()))
return loss
optimizer.step(closure)
return output_image
model = Model(device)
out_image = model.run()
out_image=out_image.cpu().clone()
denorm = transforms.Normalize((-2.12, -2.04, -1.80), (4.37, 4.46, 4.44))
out_image = out_image.squeeze()
out_image = denorm(out_image)
out_image=out_image.clamp(0,1)
unloader = transforms.ToPILImage() # reconvert into PIL image
out_image=unloader(out_image)
plt.imshow(out_image)
plt.show()
两种运行结果分别为：