-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
72 lines (63 loc) · 2.45 KB
/
model.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
import torch
import torch.nn as nn
from torch.autograd import Variable
import torchvision.datasets as dset
import torch.optim as optim
import torchvision.utils as vutils
import torchvision.transforms as transforms
import torch.nn.functional as F
# G(z)
class Generator(nn.Module):
# initializers
def __init__(self, d=128):
super(Generator, self).__init__()
self.deconv1 = nn.ConvTranspose2d(100, d*8, 4, 1, 0)
self.deconv1_bn = nn.BatchNorm2d(d*8)
self.deconv2 = nn.ConvTranspose2d(d*8, d*4, 4, 2, 1)
self.deconv2_bn = nn.BatchNorm2d(d*4)
self.deconv3 = nn.ConvTranspose2d(d*4, d*2, 4, 2, 1)
self.deconv3_bn = nn.BatchNorm2d(d*2)
self.deconv4 = nn.ConvTranspose2d(d*2, d, 4, 2, 1)
self.deconv4_bn = nn.BatchNorm2d(d)
self.deconv5 = nn.ConvTranspose2d(d, 1, 4, 2, 1)
# weight_init
def weight_init(self, mean, std):
for m in self._modules:
normal_init(self._modules[m], mean, std)
# forward method
def forward(self, input):
# x = F.relu(self.deconv1(input))
x = F.relu(self.deconv1_bn(self.deconv1(input)))
x = F.relu(self.deconv2_bn(self.deconv2(x)))
x = F.relu(self.deconv3_bn(self.deconv3(x)))
x = F.relu(self.deconv4_bn(self.deconv4(x)))
x = F.tanh(self.deconv5(x))
return x
class Discriminator(nn.Module):
# initializers
def __init__(self, d=16):
super(Discriminator, self).__init__()
self.conv1 = nn.Conv2d(1, d, 4, 2, 1)
self.conv2 = nn.Conv2d(d, d*2, 4, 2, 1)
self.conv2_bn = nn.BatchNorm2d(d*2)
self.conv3 = nn.Conv2d(d*2, d*4, 4, 2, 1)
self.conv3_bn = nn.BatchNorm2d(d*4)
self.conv4 = nn.Conv2d(d*4, d*8, 4, 2, 1)
self.conv4_bn = nn.BatchNorm2d(d*8)
self.conv5 = nn.Conv2d(d*8, 1, 4, 1, 0)
# weight_init
def weight_init(self, mean, std):
for m in self._modules:
normal_init(self._modules[m], mean, std)
# forward method
def forward(self, input):
x = F.leaky_relu(self.conv1(input), 0.2)
x = F.leaky_relu(self.conv2_bn(self.conv2(x)), 0.2)
x = F.leaky_relu(self.conv3_bn(self.conv3(x)), 0.2)
x = F.leaky_relu(self.conv4_bn(self.conv4(x)), 0.2)
x = F.sigmoid(self.conv5(x))
return x
def normal_init(m, mean, std):
if isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Conv2d):
m.weight.data.normal_(mean, std)
m.bias.data.zero_()