原图
In [1]:
from torchvision.models import vgg19
from torchvision import transforms
from PIL import Image
from matplotlib import pyplot as plt
import torch
import torch.optim as optim
import numpy as np
In [2]:
model=vgg19(pretrained=True)
In [3]:
for i in model.parameters():
i.requires_grad_(False)
In [4]:
def load_img(path, max_size=400,shape=None):
img = Image.open(path).convert('RGB')
if(max(img.size)) > max_size:size = max_size
else:size=max(img.size)
if shape is not None:
size = shape
transform = transforms.Compose([
transforms.Resize(size),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))
])
img = transform(img)[:3,:,:].unsqueeze(0)
return img
In [5]:
content=load_img('QQ图片20220207183552.jpg')
style = load_img('201609281450067826.jpg', shape=content.shape[-2:])
In [6]:
model.parameters
Out[6]:
In [7]:
def get_features(img):
features = {}
layers = {'0':'conv1_1',
'5':'conv2_1',
'10':'conv3_1',
'19':'conv4_1',
'21':'conv4_2', #content层
'28':'conv5_1'}
num=0
for name,layer in model._modules.items():
n=0
if name=='features':
for i in layer.modules():
if n==0:n+=1;continue
img=i(img)
if n==1 or n==6 or n==11 or n==20 or n==22 or n==29:
features[str(n-1)]=img
n+=1
break
return features
In [8]:
content_features=get_features(content)
style_features=get_features(style)
In [9]:
def gram_matrix(tensor):
_, d, h, w = tensor.size()
tensor = tensor.view(d, h*w)
gram = torch.mm(tensor, tensor.t())
return gram
In [10]:
style_grams = {layer:gram_matrix(style_features[layer]) for layer in style_features}
In [11]:
target = content.clone().requires_grad_(True)
In [12]:
def im_convert(tensor):
img = tensor.clone().detach()
img = img.numpy().squeeze()
img = img.transpose(1,2,0)
img = img * np.array((0.229, 0.224, 0.225)) + np.array((0.485, 0.456, 0.406))
img = img.clip(0,1)
return img
In [13]:
'''定义不同层的权重'''
style_weights = {
'0': 1,
'5': 0.8,
'10': 0.5,
'19': 0.3,
'28': 0.1,
}
'''定义2种损失对应的权重'''
content_weight = 1
style_weight = 1e6
In [14]:
show_every = 100
optimizer = optim.Adam([target], lr=0.003)
steps =5000
for ii in range(steps):
target_features = get_features(target)
content_loss = torch.mean((content_features['21'] - target_features['21'])**2)
style_loss = 0
'''加上每一层的gram_matrix矩阵的损失'''
for layer in style_weights:
target_feature = target_features[layer]
target_gram = gram_matrix(target_feature)
_, d, h, w = target_feature.shape
style_gram = style_grams[layer]
layer_style_loss = style_weights[layer] * torch.mean((target_gram - style_gram)**2)
style_loss += layer_style_loss/(d*h*w) #加到总的style_loss里，除以大小
total_loss = content_weight * content_loss + style_weight * style_loss
print('Total Loss:',total_loss.item())
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
if ii % show_every == 0 :
print('Total Loss:',total_loss.item())
plt.imshow(im_convert(target))
plt.show()
In [15]:
plt.imshow(im_convert(target))
plt.show()
