|
import torch |
|
import torch.nn as nn |
|
import torchvision.models as models |
|
|
|
|
|
class PerceptualLoss(nn.Module): |
|
r""" |
|
Perceptual loss, VGG-based |
|
https://arxiv.org/abs/1603.08155 |
|
https://github.com/dxyang/StyleTransfer/blob/master/utils.py |
|
""" |
|
|
|
def __init__(self, weights=[1.0, 1.0, 1.0, 1.0, 1.0]): |
|
super(PerceptualLoss, self).__init__() |
|
self.add_module('vgg', VGG19()) |
|
self.criterion = torch.nn.L1Loss() |
|
self.weights = weights |
|
|
|
def __call__(self, x, y): |
|
|
|
x_vgg, y_vgg = self.vgg(x), self.vgg(y) |
|
|
|
content_loss = 0.0 |
|
content_loss += self.weights[0] * self.criterion(x_vgg['relu1_1'], y_vgg['relu1_1']) |
|
content_loss += self.weights[1] * self.criterion(x_vgg['relu2_1'], y_vgg['relu2_1']) |
|
content_loss += self.weights[2] * self.criterion(x_vgg['relu3_1'], y_vgg['relu3_1']) |
|
content_loss += self.weights[3] * self.criterion(x_vgg['relu4_1'], y_vgg['relu4_1']) |
|
content_loss += self.weights[4] * self.criterion(x_vgg['relu5_1'], y_vgg['relu5_1']) |
|
|
|
|
|
return content_loss |
|
|
|
|
|
class VGG19(torch.nn.Module): |
|
def __init__(self): |
|
super(VGG19, self).__init__() |
|
features = models.vgg19(pretrained=True).features |
|
self.relu1_1 = torch.nn.Sequential() |
|
self.relu1_2 = torch.nn.Sequential() |
|
|
|
self.relu2_1 = torch.nn.Sequential() |
|
self.relu2_2 = torch.nn.Sequential() |
|
|
|
self.relu3_1 = torch.nn.Sequential() |
|
self.relu3_2 = torch.nn.Sequential() |
|
self.relu3_3 = torch.nn.Sequential() |
|
self.relu3_4 = torch.nn.Sequential() |
|
|
|
self.relu4_1 = torch.nn.Sequential() |
|
self.relu4_2 = torch.nn.Sequential() |
|
self.relu4_3 = torch.nn.Sequential() |
|
self.relu4_4 = torch.nn.Sequential() |
|
|
|
self.relu5_1 = torch.nn.Sequential() |
|
self.relu5_2 = torch.nn.Sequential() |
|
self.relu5_3 = torch.nn.Sequential() |
|
self.relu5_4 = torch.nn.Sequential() |
|
|
|
for x in range(2): |
|
self.relu1_1.add_module(str(x), features[x]) |
|
|
|
for x in range(2, 4): |
|
self.relu1_2.add_module(str(x), features[x]) |
|
|
|
for x in range(4, 7): |
|
self.relu2_1.add_module(str(x), features[x]) |
|
|
|
for x in range(7, 9): |
|
self.relu2_2.add_module(str(x), features[x]) |
|
|
|
for x in range(9, 12): |
|
self.relu3_1.add_module(str(x), features[x]) |
|
|
|
for x in range(12, 14): |
|
self.relu3_2.add_module(str(x), features[x]) |
|
|
|
for x in range(14, 16): |
|
self.relu3_2.add_module(str(x), features[x]) |
|
|
|
for x in range(16, 18): |
|
self.relu3_4.add_module(str(x), features[x]) |
|
|
|
for x in range(18, 21): |
|
self.relu4_1.add_module(str(x), features[x]) |
|
|
|
for x in range(21, 23): |
|
self.relu4_2.add_module(str(x), features[x]) |
|
|
|
for x in range(23, 25): |
|
self.relu4_3.add_module(str(x), features[x]) |
|
|
|
for x in range(25, 27): |
|
self.relu4_4.add_module(str(x), features[x]) |
|
|
|
for x in range(27, 30): |
|
self.relu5_1.add_module(str(x), features[x]) |
|
|
|
for x in range(30, 32): |
|
self.relu5_2.add_module(str(x), features[x]) |
|
|
|
for x in range(32, 34): |
|
self.relu5_3.add_module(str(x), features[x]) |
|
|
|
for x in range(34, 36): |
|
self.relu5_4.add_module(str(x), features[x]) |
|
|
|
|
|
for param in self.parameters(): |
|
param.requires_grad = False |
|
|
|
def forward(self, x): |
|
relu1_1 = self.relu1_1(x) |
|
relu1_2 = self.relu1_2(relu1_1) |
|
|
|
relu2_1 = self.relu2_1(relu1_2) |
|
relu2_2 = self.relu2_2(relu2_1) |
|
|
|
relu3_1 = self.relu3_1(relu2_2) |
|
relu3_2 = self.relu3_2(relu3_1) |
|
relu3_3 = self.relu3_3(relu3_2) |
|
relu3_4 = self.relu3_4(relu3_3) |
|
|
|
relu4_1 = self.relu4_1(relu3_4) |
|
relu4_2 = self.relu4_2(relu4_1) |
|
relu4_3 = self.relu4_3(relu4_2) |
|
relu4_4 = self.relu4_4(relu4_3) |
|
|
|
relu5_1 = self.relu5_1(relu4_4) |
|
relu5_2 = self.relu5_2(relu5_1) |
|
relu5_3 = self.relu5_3(relu5_2) |
|
relu5_4 = self.relu5_4(relu5_3) |
|
|
|
out = { |
|
'relu1_1': relu1_1, |
|
'relu1_2': relu1_2, |
|
|
|
'relu2_1': relu2_1, |
|
'relu2_2': relu2_2, |
|
|
|
'relu3_1': relu3_1, |
|
'relu3_2': relu3_2, |
|
'relu3_3': relu3_3, |
|
'relu3_4': relu3_4, |
|
|
|
'relu4_1': relu4_1, |
|
'relu4_2': relu4_2, |
|
'relu4_3': relu4_3, |
|
'relu4_4': relu4_4, |
|
|
|
'relu5_1': relu5_1, |
|
'relu5_2': relu5_2, |
|
'relu5_3': relu5_3, |
|
'relu5_4': relu5_4, |
|
} |
|
return out |
|
|
