File size: 2,596 Bytes
12b5a88 |
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 71 72 73 74 75 76 77 |
import torch
import torch.nn as nn
from adain import AdaIN
from utils import *
class StyleTransfer(nn.Module):
def __init__(self, encoder, decoder):
super(StyleTransfer, self).__init__()
layers = list(encoder.children())
self.enc_1 = nn.Sequential(*layers[:4]) # input -> relu1_1
self.enc_2 = nn.Sequential(*layers[4:11]) # relu1_1 -> relu2_1
self.enc_3 = nn.Sequential(*layers[11:18]) # relu2_1 -> relu3_1
self.enc_4 = nn.Sequential(*layers[18:31]) # relu3_1 -> relu4_1]
self.relus = [self.enc_1, self.enc_2, self.enc_3, self.enc_4]
self.decoder = decoder
self.mse = nn.MSELoss()
self.adain = AdaIN()
for name in ['enc_1', 'enc_2', 'enc_3', 'enc_4']:
for param in getattr(self, name).parameters():
param.requires_grad = False
def encode_with_save(self, input):
results = [input]
for i in range(4):
func = getattr(self, 'enc_{:d}'.format(i + 1))
results.append(func(results[-1]))
return results[1:]
def encode(self, input):
res = input
for layer in self.relus:
res = layer(res)
return res
def forward(self, content, style):
if not self.training:
self.adain.eval()
encoded_style = self.encode_with_save(style)
encoded_content = self.encode(content)
t = self.adain(encoded_content, encoded_style[-1])
g_t = self.decoder(t)
if not self.training:
return g_t
g_t_encoding = self.encode_with_save(g_t)
s_loss = self.style_loss(g_t_encoding, encoded_style)
c_loss = self.content_loss(g_t_encoding[-1], t)
return g_t, s_loss, c_loss
def style_loss(self, encoded_image, encoded_style):
MSE = torch.nn.MSELoss()
initial_mean_image, initial_std_image = mean_and_std_of_image(encoded_image[0])
initial_mean_style, initial_std_style = mean_and_std_of_image(encoded_style[0])
loss = MSE(initial_mean_image, initial_mean_style) + MSE(initial_std_image, initial_std_style)
for i in range(1, 4, 1):
mean_image, std_image = mean_and_std_of_image(encoded_image[i])
mean_style, std_style = mean_and_std_of_image(encoded_style[i])
loss += MSE(mean_image, mean_style) + MSE(std_image, std_style)
return loss
def content_loss(self, encoded_image, style_content_combined):
MSE = torch.nn.MSELoss()
return MSE(encoded_image, style_content_combined)
|