-
Notifications
You must be signed in to change notification settings - Fork 1
/
WassersteingGAN.py
102 lines (79 loc) · 2.68 KB
/
WassersteingGAN.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
## Wassersteing GAN using Saddle point Dual Averaging Optimizer in Pytorch
## Author: Amol Damare, Arjun Krishna
## SBU ID: 107914028
## contact: adamare@cs.stonybrook.edu, arjkrishna@cs.stonybrook.edu
## Reference : Yurii Nesterov. Primal-dual subgradient methods for convex problems. Mathematical Programming, 120(1):221–259, August 2009.\
## Reference. : https://github.com/wiseodd/generative-models/blob/master/GAN/wasserstein_gan/wgan_pytorch.py
import torch
import torch.nn
import torch.nn.functional as nn
import torch.autograd as autograd
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import os
from torch.autograd import Variable
from tensorflow.examples.tutorials.mnist import input_data
from Optimizer import SDAOptimizer
## Import mnist data
mnist = input_data.read_data_sets('../../MNIST_data', one_hot=True)
mb_size = 32
z_dim = 10
X_dim = mnist.train.images.shape[1]
y_dim = mnist.train.labels.shape[1]
h_dim = 128
cnt = 0
lr = 1e-4
### Generative Model
G = torch.nn.Sequential(
torch.nn.Linear(z_dim, h_dim),
torch.nn.ReLU(),
torch.nn.Linear(h_dim, X_dim),
torch.nn.Sigmoid()
)
## Descriminator
D = torch.nn.Sequential(
torch.nn.Linear(X_dim, h_dim),
torch.nn.ReLU(),
torch.nn.Linear(h_dim, 1),
)
def reset_grad():
G.zero_grad()
D.zero_grad()
## Create nestorov saddle point optimizer
solver=NestorovSaddlePointOptimizer(G.parameters(),D.parameters())
## Start training
for it in range(100000):
# Generator forward-loss-backward-update
X, _ = mnist.train.next_batch(mb_size)
X = Variable(torch.from_numpy(X))
z = Variable(torch.randn(mb_size, z_dim))
G_sample = G(z)
D_fake = D(G_sample)
D_real = D(X)
D_loss = torch.mean(D_real) - torch.mean(D_fake)
D_loss.backward()
solver.step()
# Housekeeping - reset gradient
reset_grad()
# Print and plot every now and then
if it % 1000 == 0:
print('Iter-{}; D_loss: {}; G_loss: {}'
.format(it, D_loss.data.numpy(), G_loss.data.numpy()))
samples = G(z).data.numpy()[:16]
fig = plt.figure(figsize=(4, 4))
gs = gridspec.GridSpec(4, 4)
gs.update(wspace=0.05, hspace=0.05)
for i, sample in enumerate(samples):
ax = plt.subplot(gs[i])
plt.axis('off')
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_aspect('equal')
plt.imshow(sample.reshape(28, 28), cmap='Greys_r')
if not os.path.exists('out/'):
os.makedirs('out/')
plt.savefig('out/{}.png'.format(str(cnt).zfill(3)), bbox_inches='tight')
cnt += 1
plt.close(fig)