-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
109 lines (81 loc) · 3.09 KB
/
model.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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import torch
import copy
import torch.nn.init as init
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.optim.lr_scheduler import LambdaLR
from generator import Generator
from sp_loss import CPLoss, GPLoss
USE_FLOAT16 = True
if USE_FLOAT16:
from apex import amp
class Model:
def __init__(self, device, num_steps):
# in and out channels
# for the generator:
a, b = 1, 3
def weights_init(m):
if isinstance(m, nn.Conv2d):
init.normal_(m.weight, std=0.02)
if m.bias is not None:
init.zeros_(m.bias)
elif isinstance(m, nn.BatchNorm2d):
init.ones_(m.weight)
init.zeros_(m.bias)
G = Generator(a, b).train()
self.G = G.apply(weights_init).to(device)
# it turns out that this is important
init.normal_(self.G.end[0].weight, std=1e-4)
def lambda_rule(i):
decay = num_steps // 4
m = 1.0 if i < decay else 1.0 - (i - decay) / (num_steps - decay)
return max(m, 1e-3)
self.optimizer = optim.Adam(self.G.parameters(), lr=2e-4, betas=(0.9, 0.999))
self.scheduler = LambdaLR(self.optimizer, lr_lambda=lambda_rule)
self.cp_loss = CPLoss()
self.gp_loss = GPLoss()
if USE_FLOAT16:
self.G, self.optimizer = amp.initialize(self.G, self.optimizer, opt_level='O2')
# a copy for exponential moving average
self.G_ema = copy.deepcopy(self.G)
def train_step(self, A, B):
"""
The input tensors represent images
with pixel values in [0, 1] range.
Arguments:
A: a float tensor with shape [n, a, h, w].
B: a float tensor with shape [n, b, h, w].
Returns:
a dict with float numbers.
"""
self.optimizer.zero_grad()
B_restored = self.G(A)
# it has shape [n, b, h, w]
cp_loss = self.cp_loss(B_restored, B)
gp_loss = self.gp_loss(B_restored, B)
reconstruction_loss = cp_loss + gp_loss
if not USE_FLOAT16:
reconstruction_loss.backward()
else:
with amp.scale_loss(reconstruction_loss, self.optimizer) as loss_scaled:
loss_scaled.backward()
self.optimizer.step()
self.scheduler.step()
# running average of weights
accumulate(self.G_ema, self.G)
loss_dict = {
'total_loss': reconstruction_loss.item(),
'cp_loss': cp_loss.item(),
'gp_loss': gp_loss.item()
}
return loss_dict
def save_model(self, model_path):
torch.save(self.G.state_dict(), model_path + '_generator.pth')
torch.save(self.G_ema.state_dict(), model_path + '_generator_ema.pth')
def accumulate(model_accumulator, model, decay=0.99):
"""Exponential moving average."""
params = dict(model.named_parameters())
ema_params = dict(model_accumulator.named_parameters())
for k in params.keys():
ema_params[k].data.mul_(decay).add_(1.0 - decay, params[k].data)