-
Notifications
You must be signed in to change notification settings - Fork 0
/
loss.py
70 lines (58 loc) · 2.18 KB
/
loss.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import torch
from torch import nn
from torchvision.models.vgg import vgg16, VGG16_Weights
class GeneratorLoss(nn.Module):
def __init__(self):
super(GeneratorLoss, self).__init__()
vgg = vgg16(weights=VGG16_Weights.DEFAULT)
loss_network = nn.Sequential(*list(vgg.features)[:31]).eval()
for param in loss_network.parameters():
param.requires_grad = False
self.loss_network = loss_network
self.mse_loss = nn.MSELoss()
self.tv_loss = TVLoss()
self.adversarial_loss = nn.BCELoss()
def forward(self, out_labels, out_images, target_images):
# Adversarial Loss
adversarial_loss = self.adversarial_loss(
out_labels, torch.ones_like(out_labels)
)
# Perception Loss
perception_loss = self.mse_loss(
self.loss_network(out_images), self.loss_network(target_images)
)
# Image Loss
image_loss = self.mse_loss(out_images, target_images)
# TV Loss
tv_loss = self.tv_loss(out_images)
return (
image_loss
+ 0.001 * adversarial_loss
+ 0.006 * perception_loss
+ 2e-8 * tv_loss
)
class PixelLoss(nn.Module):
def __init__(self):
super(PixelLoss, self).__init__()
self.mse_loss = nn.MSELoss()
def forward(self, out_images, target_images):
return self.mse_loss(out_images, target_images)
class TVLoss(nn.Module):
def __init__(self, tv_loss_weight=1):
super(TVLoss, self).__init__()
self.tv_loss_weight = tv_loss_weight
def forward(self, x):
batch_size = x.size()[0]
h_x = x.size()[2]
w_x = x.size()[3]
count_h = self.tensor_size(x[:, :, 1:, :])
count_w = self.tensor_size(x[:, :, :, 1:])
h_tv = torch.pow((x[:, :, 1:, :] - x[:, :, : h_x - 1, :]), 2).sum()
w_tv = torch.pow((x[:, :, :, 1:] - x[:, :, :, : w_x - 1]), 2).sum()
return self.tv_loss_weight * 2 * (h_tv / count_h + w_tv / count_w) / batch_size
@staticmethod
def tensor_size(t):
return t.size()[1] * t.size()[2] * t.size()[3]
if __name__ == "__main__":
g_loss = GeneratorLoss()
print(g_loss)