-
Notifications
You must be signed in to change notification settings - Fork 0
/
wgan.py
328 lines (253 loc) · 12.3 KB
/
wgan.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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from torch import autograd
import time as t
import matplotlib.pyplot as plt
plt.switch_backend('agg')
import os
from datasets import AudioDataset
from torchvision import utils
from resgan_model import ResGenerator2D, ResDiscriminator2D
def __init__(self, channels):
super().__init__()
self.output = nn.Tanh()
def forward(self, x):
x = self.main_module(x)
return self.output(x)
def forward(self, x):
x = self.main_module(x)
return self.output(x)
def feature_extraction(self, x):
# Use discriminator for feature extraction then flatten to vector of 16384
x = self.main_module(x)
return x.view(-1, 1024*4*4)
class WGAN_GP(object):
def __init__(self, args):
print("WGAN_GradientPenalty init model.")
self.G = ResGenerator2D
self.D = ResDiscriminator2D
self.C = args.channels
# Check if cuda is available
self.check_cuda(args.cuda)
# WGAN values from paper
self.learning_rate = 1e-4
self.b1 = 0.5
self.b2 = 0.999
self.batch_size = 64
# WGAN_gradient penalty uses ADAM
self.d_optimizer = optim.RMSProp(self.D.parameters(), lr=self.learning_rate, alpha=0.99)
self.g_optimizer = optim.RMSProp(self.G.parameters(), lr=self.learning_rate, alpha=0.99)
self.generator_iters = args.generator_iters
self.critic_iter = 5
self.lambda_term = 10
self.start_GPU = args['start_GPU']
self.device = torch.device(f"cuda:{self.start_GPU}" if (torch.cuda.is_available() and self.num_GPU > 0) else "cpu")
def train(self, train_loader):
self.t_begin = t.time()
# Now batches are callable self.data.next()
self.data = self.get_infinite_batches(train_loader)
one = torch.FloatTensor([1]).to(self.device)
mone = (one * -1).to(self.device)
for g_iter in range(self.generator_iters):
# Requires grad, Generator requires_grad = False
for p in self.D.parameters():
p.requires_grad = True
d_loss_real = 0
d_loss_fake = 0
Wasserstein_D = 0
# Train Dicriminator forward-loss-backward-update self.critic_iter times while 1 Generator forward-loss-backward-update
for d_iter in range(self.critic_iter):
self.D.zero_grad()
images = self.data.__next__()
# Check for batch to have full batch_size
if (images.size()[0] != self.batch_size):
continue
z = torch.rand((self.batch_size, 100, 1, 1))
if self.cuda:
images, z = Variable(images.cuda(self.cuda_index)), Variable(z.cuda(self.cuda_index))
else:
images, z = Variable(images), Variable(z)
# Train discriminator
# WGAN - Training discriminator more iterations than generator
# Train with real images
d_loss_real = self.D(images)
d_loss_real = d_loss_real.mean()
d_loss_real.backward(mone)
# Train with fake images
if self.cuda:
z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index)
else:
z = Variable(torch.randn(self.batch_size, 100, 1, 1))
fake_images = self.G(z)
d_loss_fake = self.D(fake_images)
d_loss_fake = d_loss_fake.mean()
d_loss_fake.backward(one)
# Train with gradient penalty
gradient_penalty = self.calculate_gradient_penalty(images.data, fake_images.data)
gradient_penalty.backward()
d_loss = d_loss_fake - d_loss_real + gradient_penalty
Wasserstein_D = d_loss_real - d_loss_fake
self.d_optimizer.step()
# Generator update
for p in self.D.parameters():
p.requires_grad = False # to avoid computation
self.G.zero_grad()
# train generator
# compute loss with fake images
z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index)
fake_images = self.G(z)
g_loss = self.D(fake_images)
g_loss = g_loss.mean()
g_loss.backward(mone)
g_cost = -g_loss
self.g_optimizer.step()
# Saving model and sampling images every 1000th generator iterations
if (g_iter) % 1000 == 0:
self.save_model()
# # Workaround because graphic card memory can't store more than 830 examples in memory for generating image
# # Therefore doing loop and generating 800 examples and stacking into list of samples to get 8000 generated images
# # This way Inception score is more correct since there are different generated examples from every class of Inception model
# sample_list = []
# for i in range(125):
# samples = self.data.__next__()
# # z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index)
# # samples = self.G(z)
# sample_list.append(samples.data.cpu().numpy())
# #
# # # Flattening list of list into one list
# new_sample_list = list(chain.from_iterable(sample_list))
# print("Calculating Inception Score over 8k generated images")
# # # Feeding list of numpy arrays
# inception_score = get_inception_score(new_sample_list, cuda=True, batch_size=32,
# resize=True, splits=10)
if not os.path.exists('training_result_images/'):
os.makedirs('training_result_images/')
# Denormalize images and save them in grid 8x8
z = Variable(torch.randn(800, 100, 1, 1)).cuda(self.cuda_index)
samples = self.G(z)
samples = samples.mul(0.5).add(0.5)
samples = samples.data.cpu()[:64]
grid = utils.make_grid(samples)
utils.save_image(grid, 'training_result_images/img_generatori_iter_{}.png'.format(str(g_iter).zfill(3)))
# Testing
time = t.time() - self.t_begin
#print("Real Inception score: {}".format(inception_score))
print("Generator iter: {}".format(g_iter))
print("Time {}".format(time))
# Write to file inception_score, gen_iters, time
#output = str(g_iter) + " " + str(time) + " " + str(inception_score[0]) + "\n"
#self.file.write(output)
# ============ TensorBoard logging ============#
# (1) Log the scalar values
info = {
'Wasserstein distance': Wasserstein_D.data[0],
'Loss D': d_loss.data[0],
'Loss G': g_cost.data[0],
'Loss D Real': d_loss_real.data[0],
'Loss D Fake': d_loss_fake.data[0]
}
for tag, value in info.items():
self.logger.scalar_summary(tag, value, g_iter + 1)
# (3) Log the images
info = {
'real_images': self.real_images(images, self.number_of_images),
'generated_images': self.generate_img(z, self.number_of_images)
}
for tag, images in info.items():
self.logger.image_summary(tag, images, g_iter + 1)
self.t_end = t.time()
print('Time of training-{}'.format((self.t_end - self.t_begin)))
#self.file.close()
# Save the trained parameters
self.save_model()
def evaluate(self, test_loader, D_model_path, G_model_path):
self.load_model(D_model_path, G_model_path)
z = Variable(torch.randn(self.batch_size, 100, 1, 1)).cuda(self.cuda_index)
samples = self.G(z)
samples = samples.mul(0.5).add(0.5)
samples = samples.data.cpu()
grid = utils.make_grid(samples)
print("Grid of 8x8 images saved to 'dgan_model_image.png'.")
utils.save_image(grid, 'dgan_model_image.png')
def calculate_gradient_penalty(self, real_images, fake_images):
eta = torch.FloatTensor(self.batch_size,1,1,1).uniform_(0,1)
eta = eta.expand(self.batch_size, real_images.size(1), real_images.size(2), real_images.size(3))
if self.cuda:
eta = eta.cuda(self.cuda_index)
else:
eta = eta
interpolated = eta * real_images + ((1 - eta) * fake_images)
if self.cuda:
interpolated = interpolated.cuda(self.cuda_index)
else:
interpolated = interpolated
# define it to calculate gradient
interpolated = Variable(interpolated, requires_grad=True)
# calculate probability of interpolated examples
prob_interpolated = self.D(interpolated)
# calculate gradients of probabilities with respect to examples
gradients = autograd.grad(outputs=prob_interpolated, inputs=interpolated,
grad_outputs=torch.ones(
prob_interpolated.size()).cuda(self.cuda_index) if self.cuda else torch.ones(
prob_interpolated.size()),
create_graph=True, retain_graph=True)[0]
grad_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * self.lambda_term
return grad_penalty
def real_images(self, images, number_of_images):
if (self.C == 3):
return self.to_np(images.view(-1, self.C, 32, 32)[:self.number_of_images])
else:
return self.to_np(images.view(-1, 32, 32)[:self.number_of_images])
def generate_img(self, z, number_of_images):
samples = self.G(z).data.cpu().numpy()[:number_of_images]
generated_images = []
for sample in samples:
if self.C == 3:
generated_images.append(sample.reshape(self.C, 32, 32))
else:
generated_images.append(sample.reshape(32, 32))
return generated_images
def to_np(self, x):
return x.data.cpu().numpy()
def save_model(self):
torch.save(self.G.state_dict(), './generator.pkl')
torch.save(self.D.state_dict(), './discriminator.pkl')
print('Models save to ./generator.pkl & ./discriminator.pkl ')
def load_model(self, D_model_filename, G_model_filename):
D_model_path = os.path.join(os.getcwd(), D_model_filename)
G_model_path = os.path.join(os.getcwd(), G_model_filename)
self.D.load_state_dict(torch.load(D_model_path))
self.G.load_state_dict(torch.load(G_model_path))
print('Generator model loaded from {}.'.format(G_model_path))
print('Discriminator model loaded from {}-'.format(D_model_path))
def get_infinite_batches(self, data_loader):
while True:
for i, (images, _) in enumerate(data_loader):
yield images
def generate_latent_walk(self, number):
if not os.path.exists('interpolated_images/'):
os.makedirs('interpolated_images/')
number_int = 10
# interpolate between twe noise(z1, z2).
z_intp = torch.FloatTensor(1, 100, 1, 1)
z1 = torch.randn(1, 100, 1, 1)
z2 = torch.randn(1, 100, 1, 1)
if self.cuda:
z_intp = z_intp.cuda()
z1 = z1.cuda()
z2 = z2.cuda()
z_intp = Variable(z_intp)
images = []
alpha = 1.0 / float(number_int + 1)
print(alpha)
for i in range(1, number_int + 1):
z_intp.data = z1*alpha + z2*(1.0 - alpha)
alpha += alpha
fake_im = self.G(z_intp)
fake_im = fake_im.mul(0.5).add(0.5) #denormalize
images.append(fake_im.view(self.C,32,32).data.cpu())
grid = utils.make_grid(images, nrow=number_int )
utils.save_image(grid, 'interpolated_images/interpolated_{}.png'.format(str(number).zfill(3)))
print("Saved interpolated images.")