python
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import models, transforms
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
def get_model_and_losses(self):
model = nn.Sequential()
i = 0
for layer in vgg.children():
if isinstance(layer, nn.Conv2d):
i += 1
name = 'conv_{}'.format(i)
elif isinstance(layer, nn.ReLU):
name = 'relu_{}'.format(i)
layer = nn.ReLU(inplace=False)
elif isinstance(layer, nn.MaxPool2d):
name = 'pool_{}'.format(i)
elif isinstance(layer, nn.BatchNorm2d):
name = 'bn_{}'.format(i)
else:
raise RuntimeError('Unrecognized layer: {}'.format(layer.__class__.__name__))
model.add_module(name, layer)
if name == 'conv_4_2': 内容损失层
target = model(self.content_image).detach()
content_loss = ContentLoss(target)
model.add_module(content_loss_{}.format(i), content_loss)
self.content_losses.append(content_loss)
if name in ['conv_1_1', 'conv_2_1', 'conv_3_1', 'conv_4_1', 'conv_5_1']: 风格损失层
target_feature = model(self.style_image).detach()
style_loss = StyleLoss(target_feature)
model.add_module(style_loss_{}.format(i), style_loss)
self.style_losses.append(style_loss)
num_steps = 300
for i in range(num_steps):
def closure():
optimizer.zero_grad()
model.forward()
style_score = 0
content_score = 0
for sl in model.style_losses:
style_score += sl.loss
for cl in model.content_losses:
content_score += cl.loss
style_score *= 1000000
content_score *= 1
loss = style_score + content_score
loss.backward()
if i % 50 == 0:
print('Step {}: Style Loss : {:4f} Content Loss: {:4f}'.format(i, style_score.item(), content_score.item()))