-
Notifications
You must be signed in to change notification settings - Fork 0
/
gan.py
177 lines (152 loc) · 6.57 KB
/
gan.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
"""
GAN.
Everything in one file.
"""
##################################################################################################################################
import torch
import torch.nn as nn
class Generator(nn.Module):
def __init__(self, in_chan=100, hidden=64, out_chan=1):
super(Generator, self).__init__()
self.decoder = nn.Sequential(
nn.ConvTranspose2d(in_chan, hidden * 8, 4, 1, 0), # 4x4
nn.BatchNorm2d(hidden * 8),
nn.ReLU(True),
nn.ConvTranspose2d(hidden * 8, hidden * 4, 4, 2, 1), # 8x8
nn.BatchNorm2d(hidden * 4),
nn.ReLU(True),
nn.ConvTranspose2d(hidden * 4, hidden * 2, 4, 2, 1), # 16x16
nn.BatchNorm2d(hidden * 2),
nn.ReLU(True),
nn.ConvTranspose2d(hidden * 2, hidden, 4, 2, 1), # 32x32
nn.BatchNorm2d(hidden),
nn.ReLU(True),
nn.ConvTranspose2d(hidden, out_chan, kernel_size=1, stride=1, padding=2), # 28x28
nn.Tanh() # [-1, 1]
)
def forward(self, x):
output = self.decoder(x)
return output
class Discriminator(nn.Module):
def __init__(self, in_chan=1, hidden=64):
super(Discriminator, self).__init__()
self.classifier = nn.Sequential(
nn.Conv2d(in_chan, hidden, 4, 2, 1), # H/2, W/2
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(hidden, hidden * 2, 4, 2, 1), # H/4, W/4
nn.BatchNorm2d(hidden * 2),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(hidden * 2, hidden * 4, 4, 2, 1), # H/8, W/8
nn.BatchNorm2d(hidden * 4),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(hidden * 4, 1, 4, 2, 1), # H/16, W/16
nn.Sigmoid()
)
def forward(self, x):
output = self.classifier(x)
return output.view(-1, 1).squeeze(1)
##################################################################################################################################
import torchvision
from tqdm.auto import tqdm
def get_dataloader(batch_size=128):
transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
dataset = torchvision.datasets.mnist.MNIST("./data", download=True, train=True, transform=transform)
return torch.utils.data.DataLoader(dataset, batch_size, shuffle=True)
def get_device():
device = 'cpu'
if torch.backends.mps.is_available():
device = 'mps:0'
if torch.cuda.is_available():
device = 'cuda'
return device
def train(n_epochs, batch_size=128, latent=8, hidden=16):
device = get_device()
dataloader = get_dataloader(batch_size=batch_size)
netG = Generator(latent, hidden, 1).to(device)
netD = Discriminator(1, hidden).to(device)
criterion = nn.BCELoss()
optimizerD = torch.optim.Adam(netD.parameters())
optimizerG = torch.optim.Adam(netG.parameters())
netG.train()
netD.train()
with tqdm(range(n_epochs), colour="#00ee00") as epoch_pbar:
for _ in epoch_pbar:
with tqdm(dataloader, leave=False, colour="#005500") as batch_pbar:
for images, _ in batch_pbar:
############################
batch_size = images.size(0)
positive_labels = torch.full((batch_size,), 1.0, device=device)
negative_labels = torch.full((batch_size,), 0.0, device=device)
real_image = images.to(device)
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
netD.zero_grad()
# train with real for D
output_D = netD(real_image)
loss_real = criterion(output_D, positive_labels)
loss_real.backward()
# train with fake for G
noise_latent = torch.randn(batch_size, latent, 1, 1, device=device)
fake_images = netG(noise_latent)
output_D = netD(fake_images.detach())
loss_fake = criterion(output_D, negative_labels)
loss_fake.backward()
# update weights for D
optimizerD.step()
# (2) Update G network: maximize log(D(G(z)))
netG.zero_grad()
output_D = netD(fake_images) # attach for training G
loss_g = criterion(output_D, positive_labels) # train G to let it looks like real
loss_g.backward()
optimizerG.step()
batch_pbar.set_description(f'{loss_real.item():.3f}, {loss_fake.item():.3f}, {loss_g.item():.3f}')
epoch_pbar.set_description(f'{loss_real.item():.3f}, {loss_fake.item():.3f}, {loss_g.item():.3f}')
torch.save(netG.state_dict(), 'gan_g.pth')
torch.save(netD.state_dict(), 'gan_d.pth')
##################################################################################################################################
import matplotlib.pyplot as plt
def show_images(images):
# Converting images to CPU numpy arrays
if type(images) is torch.Tensor:
images = images.detach().cpu().numpy()
# Defining number of rows and columns
fig = plt.figure(figsize=(4, 4))
rows = int(len(images) ** (1 / 2))
cols = round(len(images) / rows)
# Populating figure with sub-plots
idx = 0
for r in range(rows):
for c in range(cols):
if idx < len(images):
fig.add_subplot(rows, cols, idx + 1)
plt.imshow(images[idx][0], cmap="gray")
plt.axis('off')
idx += 1
plt.tight_layout()
plt.show()
def predict(n_samples=64, latent=8, hidden=16):
device = get_device()
net = Generator(latent, hidden, 1).to(device)
net.load_state_dict(torch.load('gan_g.pth'))
net.eval()
with torch.no_grad():
z = torch.randn(n_samples, latent, 1, 1).to(device)
x = net(z)
show_images(x)
##################################################################################################################################
from absl import flags
from absl import app
def main(unused_args):
"""
Samples:
python gan.py --train --epochs 100 --predict
"""
if FLAGS.train:
train(n_epochs=FLAGS.epochs)
if FLAGS.predict:
predict()
if __name__ == '__main__':
FLAGS = flags.FLAGS
flags.DEFINE_bool("train", False, "Train the model")
flags.DEFINE_bool("predict", False, "Predict")
flags.DEFINE_integer("epochs", 3, "Epochs to train")
app.run(main)