-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathvae.py
103 lines (80 loc) · 3.45 KB
/
vae.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
import torch
import torch.utils.data
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
class VAE(nn.Module):
def __init__(self, nc, ngf, ndf, latent_variable_size):
super(VAE, self).__init__()
self.nc = nc
self.ngf = ngf
self.ndf = ndf
self.latent_variable_size = latent_variable_size
# encoder
self.e1 = nn.Conv2d(nc, ndf, 4, 2, 1)
self.bn1 = nn.BatchNorm2d(ndf)
self.e2 = nn.Conv2d(ndf, ndf*2, 4, 2, 1)
self.bn2 = nn.BatchNorm2d(ndf*2)
self.e3 = nn.Conv2d(ndf*2, ndf*4, 4, 2, 1)
self.bn3 = nn.BatchNorm2d(ndf*4)
self.e4 = nn.Conv2d(ndf*4, ndf*8, 4, 2, 1)
self.bn4 = nn.BatchNorm2d(ndf*8)
self.e5 = nn.Conv2d(ndf*8, ndf*8, 4, 2, 1)
self.bn5 = nn.BatchNorm2d(ndf*8)
self.fc1 = nn.Linear(ndf*8*4*4, latent_variable_size)
self.fc2 = nn.Linear(ndf*8*4*4, latent_variable_size)
# decoder
self.d1 = nn.Linear(latent_variable_size, ngf*8*2*4*4)
self.up1 = nn.UpsamplingNearest2d(scale_factor=2)
self.pd1 = nn.ReplicationPad2d(1)
self.d2 = nn.Conv2d(ngf*8*2, ngf*8, 3, 1)
self.bn6 = nn.BatchNorm2d(ngf*8, 1.e-3)
self.up2 = nn.UpsamplingNearest2d(scale_factor=2)
self.pd2 = nn.ReplicationPad2d(1)
self.d3 = nn.Conv2d(ngf*8, ngf*4, 3, 1)
self.bn7 = nn.BatchNorm2d(ngf*4, 1.e-3)
self.up3 = nn.UpsamplingNearest2d(scale_factor=2)
self.pd3 = nn.ReplicationPad2d(1)
self.d4 = nn.Conv2d(ngf*4, ngf*2, 3, 1)
self.bn8 = nn.BatchNorm2d(ngf*2, 1.e-3)
self.up4 = nn.UpsamplingNearest2d(scale_factor=2)
self.pd4 = nn.ReplicationPad2d(1)
self.d5 = nn.Conv2d(ngf*2, ngf, 3, 1)
self.bn9 = nn.BatchNorm2d(ngf, 1.e-3)
self.up5 = nn.UpsamplingNearest2d(scale_factor=2)
self.pd5 = nn.ReplicationPad2d(1)
self.d6 = nn.Conv2d(ngf, nc, 3, 1)
self.leakyrelu = nn.LeakyReLU(0.2)
self.relu = nn.ReLU()
self.sigmoid = nn.Sigmoid()
def encode(self, x):
h1 = self.leakyrelu(self.bn1(self.e1(x)))
h2 = self.leakyrelu(self.bn2(self.e2(h1)))
h3 = self.leakyrelu(self.bn3(self.e3(h2)))
h4 = self.leakyrelu(self.bn4(self.e4(h3)))
h5 = self.leakyrelu(self.bn5(self.e5(h4)))
h5 = h5.view(-1, self.ndf*8*4*4)
return self.fc1(h5), self.fc2(h5)
def reparametrize(self, mu, logvar):
std = logvar.mul(0.5).exp_()
eps = torch.FloatTensor(std.size()).normal_()
eps = Variable(eps)
return eps.mul(std).add_(mu)
def decode(self, z):
h1 = self.relu(self.d1(z))
h1 = h1.view(-1, self.ngf*8*2, 4, 4)
h2 = self.leakyrelu(self.bn6(self.d2(self.pd1(self.up1(h1)))))
h3 = self.leakyrelu(self.bn7(self.d3(self.pd2(self.up2(h2)))))
h4 = self.leakyrelu(self.bn8(self.d4(self.pd3(self.up3(h3)))))
h5 = self.leakyrelu(self.bn9(self.d5(self.pd4(self.up4(h4)))))
return self.sigmoid(self.d6(self.pd5(self.up5(h5))))
def get_latent_var(self, x):
mu, logvar = self.encode(x.view(-1, self.nc, self.ndf, self.ngf))
z = self.reparametrize(mu, logvar)
return z
def forward(self, x):
mu, logvar = self.encode(x.view(-1, self.nc, self.ndf, self.ngf))
z = self.reparametrize(mu, logvar)
res = self.decode(z)
return res, mu, logvar
model = VAE(nc=3, ngf=128, ndf=128, latent_variable_size=500)