-
Notifications
You must be signed in to change notification settings - Fork 0
/
gan_mnist.py
141 lines (107 loc) · 4.4 KB
/
gan_mnist.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
# prerequisites
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
from torchvision.utils import save_image
# Device configuration
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
bs = 100
# MNIST Dataset
# transform = transforms.Compose([
# transforms.ToTensor(),
# transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])
transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
train_dataset = datasets.MNIST(root='./mnist_data/', train=True, transform=transform, download=True)
test_dataset = datasets.MNIST(root='./mnist_data/', train=False, transform=transform, download=False)
# Data Loader (Input Pipeline)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=bs, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=bs, shuffle=False)
class Generator(nn.Module):
def __init__(self, g_input_dim, g_output_dim):
super(Generator, self).__init__()
self.fc1 = nn.Linear(g_input_dim, 256)
self.fc2 = nn.Linear(self.fc1.out_features, self.fc1.out_features*2)
self.fc3 = nn.Linear(self.fc2.out_features, self.fc2.out_features*2)
self.fc4 = nn.Linear(self.fc3.out_features, g_output_dim)
# forward method
def forward(self, x):
x = F.leaky_relu(self.fc1(x), 0.2)
x = F.leaky_relu(self.fc2(x), 0.2)
x = F.leaky_relu(self.fc3(x), 0.2)
return torch.tanh(self.fc4(x))
class Discriminator(nn.Module):
def __init__(self, d_input_dim):
super(Discriminator, self).__init__()
self.fc1 = nn.Linear(d_input_dim, 1024)
self.fc2 = nn.Linear(self.fc1.out_features, self.fc1.out_features//2)
self.fc3 = nn.Linear(self.fc2.out_features, self.fc2.out_features//2)
self.fc4 = nn.Linear(self.fc3.out_features, 1)
# forward method
def forward(self, x):
x = F.leaky_relu(self.fc1(x), 0.2)
x = F.dropout(x, 0.3)
x = F.leaky_relu(self.fc2(x), 0.2)
x = F.dropout(x, 0.3)
x = F.leaky_relu(self.fc3(x), 0.2)
x = F.dropout(x, 0.3)
return torch.sigmoid(self.fc4(x))
# build network
z_dim = 100
mnist_dim = train_dataset.train_data.size(1) * train_dataset.train_data.size(2)
G = Generator(g_input_dim = z_dim, g_output_dim = mnist_dim).to(device)
D = Discriminator(mnist_dim).to(device)
# loss
criterion = nn.BCELoss()
# optimizer
lr = 0.0002
G_optimizer = optim.Adam(G.parameters(), lr = lr)
D_optimizer = optim.Adam(D.parameters(), lr = lr)
def D_train(x):
#=======================Train the discriminator=======================#
D.zero_grad()
# train discriminator on real
x_real, y_real = x.view(-1, mnist_dim), torch.ones(bs, 1)
x_real, y_real = Variable(x_real.to(device)), Variable(y_real.to(device))
D_output = D(x_real)
D_real_loss = criterion(D_output, y_real)
D_real_score = D_output
# train discriminator on facke
z = Variable(torch.randn(bs, z_dim).to(device))
x_fake, y_fake = G(z), Variable(torch.zeros(bs, 1).to(device))
D_output = D(x_fake)
D_fake_loss = criterion(D_output, y_fake)
D_fake_score = D_output
# gradient backprop & optimize ONLY D's parameters
D_loss = D_real_loss + D_fake_loss
D_loss.backward()
D_optimizer.step()
return D_loss.data.item()
def G_train(x):
#=======================Train the generator=======================#
G.zero_grad()
z = Variable(torch.randn(bs, z_dim).to(device))
y = Variable(torch.ones(bs, 1).to(device))
G_output = G(z)
D_output = D(G_output)
G_loss = criterion(D_output, y)
# gradient backprop & optimize ONLY G's parameters
G_loss.backward()
G_optimizer.step()
return G_loss.data.item()
n_epoch = 200
for epoch in range(1, n_epoch+1):
D_losses, G_losses = [], []
for batch_idx, (x, _) in enumerate(train_loader):
D_losses.append(D_train(x))
G_losses.append(G_train(x))
print('[%d/%d]: loss_d: %.3f, loss_g: %.3f' % (
(epoch), n_epoch, torch.mean(torch.FloatTensor(D_losses)), torch.mean(torch.FloatTensor(G_losses))))
with torch.no_grad():
test_z = Variable(torch.randn(bs, z_dim).to(device))
generated = G(test_z)
save_image(generated.view(generated.size(0), 1, 28, 28), './drive/MyDrive/samples/sample_' + '.png')