Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,214 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import torch
|
3 |
+
from IPython.display import Image
|
4 |
+
from torchvision.utils import save_image
|
5 |
+
import torchvision
|
6 |
+
from torchvision.transforms import ToTensor, Normalize, Compose
|
7 |
+
from torchvision.datasets import MNIST
|
8 |
+
from torch.utils.data import DataLoader
|
9 |
+
import matplotlib.pyplot as plt
|
10 |
+
import torch.nn as nn
|
11 |
+
import cv2
|
12 |
+
import os
|
13 |
+
from IPython.display import FileLink
|
14 |
+
%matplotlib inline
|
15 |
+
|
16 |
+
mnist = MNIST(root='data',
|
17 |
+
train=True,
|
18 |
+
download=True,
|
19 |
+
transform=Compose([ToTensor(), Normalize(mean=(0.5,), std=(0.5,))]))
|
20 |
+
img, label = mnist[0]
|
21 |
+
print('Label: ', label)
|
22 |
+
print(img[:,10:15,10:15])
|
23 |
+
torch.min(img), torch.max(img)
|
24 |
+
def denorm(x):
|
25 |
+
out = (x + 1) / 2
|
26 |
+
return out.clamp(0, 1)
|
27 |
+
|
28 |
+
img_norm = denorm(img)
|
29 |
+
plt.imshow(img_norm[0], cmap='gray')
|
30 |
+
print('Label:', label)
|
31 |
+
|
32 |
+
batch_size = 100
|
33 |
+
data_loader = DataLoader(mnist, batch_size, shuffle=True)
|
34 |
+
|
35 |
+
for img_batch, label_batch in data_loader:
|
36 |
+
print('first batch')
|
37 |
+
print(img_batch.shape)
|
38 |
+
plt.imshow(img_batch[0][0], cmap='gray')
|
39 |
+
print(label_batch)
|
40 |
+
break
|
41 |
+
|
42 |
+
# Device configuration
|
43 |
+
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
44 |
+
|
45 |
+
image_size = 784
|
46 |
+
hidden_size = 256
|
47 |
+
|
48 |
+
D = nn.Sequential(
|
49 |
+
nn.Linear(image_size, hidden_size),
|
50 |
+
nn.LeakyReLU(0.2),
|
51 |
+
nn.Linear(hidden_size, hidden_size),
|
52 |
+
nn.LeakyReLU(0.2),
|
53 |
+
nn.Linear(hidden_size, 1),
|
54 |
+
nn.Sigmoid())
|
55 |
+
|
56 |
+
D.to(device);
|
57 |
+
|
58 |
+
latent_size = 64
|
59 |
+
|
60 |
+
G = nn.Sequential(
|
61 |
+
nn.Linear(latent_size, hidden_size),
|
62 |
+
nn.ReLU(),
|
63 |
+
nn.Linear(hidden_size, hidden_size),
|
64 |
+
nn.ReLU(),
|
65 |
+
nn.Linear(hidden_size, image_size),
|
66 |
+
nn.Tanh())
|
67 |
+
|
68 |
+
y = G(torch.randn(2, latent_size))
|
69 |
+
gen_imgs = denorm(y.reshape((-1, 28,28)).detach())
|
70 |
+
|
71 |
+
plt.imshow(gen_imgs[0], cmap='gray');
|
72 |
+
|
73 |
+
plt.imshow(gen_imgs[1], cmap='gray');
|
74 |
+
|
75 |
+
G.to(device);
|
76 |
+
|
77 |
+
criterion = nn.BCELoss()
|
78 |
+
d_optimizer = torch.optim.Adam(D.parameters(), lr=0.0002)
|
79 |
+
|
80 |
+
def reset_grad():
|
81 |
+
d_optimizer.zero_grad()
|
82 |
+
g_optimizer.zero_grad()
|
83 |
+
|
84 |
+
def train_discriminator(images):
|
85 |
+
# Create the labels which are later used as input for the BCE loss
|
86 |
+
real_labels = torch.ones(batch_size, 1).to(device)
|
87 |
+
fake_labels = torch.zeros(batch_size, 1).to(device)
|
88 |
+
|
89 |
+
# Loss for real images
|
90 |
+
outputs = D(images)
|
91 |
+
d_loss_real = criterion(outputs, real_labels)
|
92 |
+
real_score = outputs
|
93 |
+
|
94 |
+
# Loss for fake images
|
95 |
+
z = torch.randn(batch_size, latent_size).to(device)
|
96 |
+
fake_images = G(z)
|
97 |
+
outputs = D(fake_images)
|
98 |
+
d_loss_fake = criterion(outputs, fake_labels)
|
99 |
+
fake_score = outputs
|
100 |
+
|
101 |
+
# Combine losses
|
102 |
+
d_loss = d_loss_real + d_loss_fake
|
103 |
+
# Reset gradients
|
104 |
+
reset_grad()
|
105 |
+
# Compute gradients
|
106 |
+
d_loss.backward()
|
107 |
+
# Adjust the parameters using backprop
|
108 |
+
d_optimizer.step()
|
109 |
+
|
110 |
+
return d_loss, real_score, fake_score
|
111 |
+
|
112 |
+
g_optimizer = torch.optim.Adam(G.parameters(), lr=0.0002)
|
113 |
+
|
114 |
+
def train_generator():
|
115 |
+
# Generate fake images and calculate loss
|
116 |
+
z = torch.randn(batch_size, latent_size).to(device)
|
117 |
+
fake_images = G(z)
|
118 |
+
labels = torch.ones(batch_size, 1).to(device)
|
119 |
+
g_loss = criterion(D(fake_images), labels)
|
120 |
+
|
121 |
+
# Backprop and optimize
|
122 |
+
reset_grad()
|
123 |
+
g_loss.backward()
|
124 |
+
g_optimizer.step()
|
125 |
+
return g_loss, fake_images
|
126 |
+
|
127 |
+
sample_dir = 'samples'
|
128 |
+
if not os.path.exists(sample_dir):
|
129 |
+
os.makedirs(sample_dir)
|
130 |
+
|
131 |
+
# Save some real images
|
132 |
+
for images, _ in data_loader:
|
133 |
+
images = images.reshape(images.size(0), 1, 28, 28)
|
134 |
+
save_image(denorm(images), os.path.join(sample_dir, 'real_images.png'), nrow=10)
|
135 |
+
break
|
136 |
+
|
137 |
+
Image(os.path.join(sample_dir, 'real_images.png'))
|
138 |
+
|
139 |
+
sample_vectors = torch.randn(batch_size, latent_size).to(device)
|
140 |
+
|
141 |
+
def save_fake_images(index):
|
142 |
+
fake_images = G(sample_vectors)
|
143 |
+
fake_images = fake_images.reshape(fake_images.size(0), 1, 28, 28)
|
144 |
+
fake_fname = 'fake_images-{0:0=4d}.png'.format(index)
|
145 |
+
print('Saving', fake_fname)
|
146 |
+
save_image(denorm(fake_images), os.path.join(sample_dir, fake_fname), nrow=10)
|
147 |
+
|
148 |
+
# Before training
|
149 |
+
save_fake_images(0)
|
150 |
+
Image(os.path.join(sample_dir, 'fake_images-0000.png'))
|
151 |
+
|
152 |
+
%%time
|
153 |
+
|
154 |
+
num_epochs = 300
|
155 |
+
total_step = len(data_loader)
|
156 |
+
d_losses, g_losses, real_scores, fake_scores = [], [], [], []
|
157 |
+
|
158 |
+
for epoch in range(num_epochs):
|
159 |
+
for i, (images, _) in enumerate(data_loader):
|
160 |
+
# Load a batch & transform to vectors
|
161 |
+
images = images.reshape(batch_size, -1).to(device)
|
162 |
+
|
163 |
+
# Train the discriminator and generator
|
164 |
+
d_loss, real_score, fake_score = train_discriminator(images)
|
165 |
+
g_loss, fake_images = train_generator()
|
166 |
+
|
167 |
+
# Inspect the losses
|
168 |
+
if (i+1) % 200 == 0:
|
169 |
+
d_losses.append(d_loss.item())
|
170 |
+
g_losses.append(g_loss.item())
|
171 |
+
real_scores.append(real_score.mean().item())
|
172 |
+
fake_scores.append(fake_score.mean().item())
|
173 |
+
print('Epoch [{}/{}], Step [{}/{}], d_loss: {:.4f}, g_loss: {:.4f}, D(x): {:.2f}, D(G(z)): {:.2f}'
|
174 |
+
.format(epoch, num_epochs, i+1, total_step, d_loss.item(), g_loss.item(),
|
175 |
+
real_score.mean().item(), fake_score.mean().item()))
|
176 |
+
|
177 |
+
# Sample and save images
|
178 |
+
save_fake_images(epoch+1)
|
179 |
+
|
180 |
+
# Save the model checkpoints
|
181 |
+
torch.save(G.state_dict(), 'G.ckpt')
|
182 |
+
torch.save(D.state_dict(), 'D.ckpt')
|
183 |
+
|
184 |
+
Image('./samples/fake_images-0010.png')
|
185 |
+
|
186 |
+
Image('./samples/fake_images-0050.png')
|
187 |
+
|
188 |
+
Image('./samples/fake_images-0100.png')
|
189 |
+
|
190 |
+
Image('./samples/fake_images-0300.png')
|
191 |
+
|
192 |
+
vid_fname = 'gans_training.avi'
|
193 |
+
|
194 |
+
files = [os.path.join(sample_dir, f) for f in os.listdir(sample_dir) if 'fake_images' in f]
|
195 |
+
files.sort()
|
196 |
+
|
197 |
+
out = cv2.VideoWriter(vid_fname,cv2.VideoWriter_fourcc(*'MP4V'), 8, (302,302))
|
198 |
+
[out.write(cv2.imread(fname)) for fname in files]
|
199 |
+
out.release()
|
200 |
+
FileLink('gans_training.avi')
|
201 |
+
|
202 |
+
plt.plot(d_losses, '-')
|
203 |
+
plt.plot(g_losses, '-')
|
204 |
+
plt.xlabel('epoch')
|
205 |
+
plt.ylabel('loss')
|
206 |
+
plt.legend(['Discriminator', 'Generator'])
|
207 |
+
plt.title('Losses');
|
208 |
+
|
209 |
+
plt.plot(real_scores, '-')
|
210 |
+
plt.plot(fake_scores, '-')
|
211 |
+
plt.xlabel('epoch')
|
212 |
+
plt.ylabel('score')
|
213 |
+
plt.legend(['Real Score', 'Fake score'])
|
214 |
+
plt.title('Scores');
|