-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.lua
145 lines (128 loc) · 7.17 KB
/
model.lua
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
require 'cutorch';
require 'cunn';
require 'torch';
require 'nn';
--------------------------------------------------------------------------
-- Model --
--------------------------------------------------------------------------
encoder = nn.Sequential()
encoder:add(nn.SpatialConvolution(3, 64, 4, 4, 2, 2, 1, 1)) -- 3x128x128 -> 64x64x64
--encoder.modules[#encoder.modules].weight:normal(0, 0.02)
--encoder.modules[#encoder.modules].bias:fill(0)
encoder:add(nn.SpatialBatchNormalization(64))
-- encoder.modules[#encoder.modules].weight:normal(1.0, 0.02)
-- encoder.modules[#encoder.modules].bias:fill(0)
encoder:add(nn.ReLU(true)) --relu1
encoder:add(nn.SpatialDropout(0.25))
encoder:add(nn.SpatialConvolution(64, 64, 4, 4, 2, 2, 1, 1)) --64x64x64 -> 64x32x32
-- encoder.modules[#encoder.modules].weight:normal(0, 0.02)
-- encoder.modules[#encoder.modules].bias:fill(0.1)
encoder:add(nn.SpatialBatchNormalization(64))
-- encoder.modules[#encoder.modules].weight:normal(1.0, 0.02)
-- encoder.modules[#encoder.modules].bias:fill(0)
encoder:add(nn.ReLU(true)) --relu2
encoder:add(nn.SpatialConvolution(64, 128, 4, 4, 2, 2, 1, 1)) --64x32x32 -> 128x16x16
-- encoder.modules[#encoder.modules].weight:normal(0, 0.02)
-- encoder.modules[#encoder.modules].bias:fill(0)
encoder:add(nn.SpatialBatchNormalization(128))
-- encoder.modules[#encoder.modules].weight:normal(1.0, 0.02)
-- encoder.modules[#encoder.modules].bias:fill(0)
encoder:add(nn.ReLU(true)) --relu3
encoder:add(nn.SpatialConvolution(128, 256, 4, 4, 2, 2, 1, 1)) --128x16x16 -> 256x8x8
-- encoder.modules[#encoder.modules].weight:normal(0, 0.02)
-- encoder.modules[#encoder.modules].bias:fill(0.1)
encoder:add(nn.SpatialBatchNormalization(256))
-- encoder.modules[#encoder.modules].weight:normal(1.0, 0.02)
-- encoder.modules[#encoder.modules].bias:fill(0)
encoder:add(nn.ReLU(true)) --relu4
encoder:add(nn.SpatialConvolution(256, 512, 4, 4, 2, 2, 1, 1)) --256x8x8 -> 512x4x4
-- encoder.modules[#encoder.modules].weight:normal(0, 0.02)
-- encoder.modules[#encoder.modules].bias:fill(0.1)
encoder:add(nn.SpatialBatchNormalization(512))
-- encoder.modules[#encoder.modules].weight:normal(1.0, 0.02)
-- encoder.modules[#encoder.modules].bias:fill(0)
encoder:add(nn.ReLU(true)) --relu5
encoder:add(nn.View( 512*4*4 ))
encoder:add( nn.Linear( 512*4*4, 4500 ) ) --512*4*4 -> 4000
encoder:add( nn.Tanh() )
------------------------------------------------------------------
-- classifier for pre-training --
------------------------------------------------------------------
classifier = nn.Sequential()
classifier:add( encoder ) -- 3x128x128 -> 4000
classifier:add( nn.Linear( 4500, 400 ) )
classifier:add( nn.Tanh() )
classifier:add( nn.Linear( 400, 1 ) )
classifier:add( nn.Sigmoid() )
------------------------------------------------------------------
-- Decoder
decoder = nn.Sequential()
decoder:add( nn.Linear( 4500, 512*4*4 ) )
decoder:add( nn.Reshape( 512, 4, 4 ) )
decoder:add( nn.SpatialFullConvolution( 512, 256, 4, 4, 2, 2, 1, 1 ) ) --512x4x4 -> 256x8x8
-- decoder.modules[#decoder.modules].weight:normal(0, 0.02)
-- decoder.modules[#decoder.modules].bias:fill(0.1)
decoder:add( nn.SpatialBatchNormalization( 256 ) )
-- decoder.modules[#decoder.modules].weight:normal(1.0, 0.02)
-- decoder.modules[#decoder.modules].bias:fill(0)
decoder:add(nn.ReLU(true))
decoder:add( nn.SpatialFullConvolution( 256, 128, 4, 4, 2, 2, 1, 1 ) ) --256x8x8 -> 128x16x16
-- decoder.modules[#decoder.modules].weight:normal(0, 0.02)
-- decoder.modules[#decoder.modules].bias:fill(0.1)
decoder:add( nn.SpatialBatchNormalization( 128 ) )
-- decoder.modules[#decoder.modules].weight:normal(1.0, 0.02)
-- decoder.modules[#decoder.modules].bias:fill(0)
decoder:add(nn.ReLU(true))
decoder:add( nn.SpatialFullConvolution( 128, 64, 4, 4, 2, 2, 1, 1 ) ) --128x16x16 -> 64x32x32
-- decoder.modules[#decoder.modules].weight:normal(0, 0.02)
-- decoder.modules[#decoder.modules].bias:fill(0.1)
decoder:add( nn.SpatialBatchNormalization( 64 ) )
-- decoder.modules[#decoder.modules].weight:normal(1.0, 0.02)
-- decoder.modules[#decoder.modules].bias:fill(0)
decoder:add(nn.ReLU(true))
decoder:add( nn.SpatialFullConvolution( 64, 3, 4, 4, 2, 2, 1, 1 ) ) --64x32x32 -> 3x64x64
decoder:add(nn.Tanh())
generator = nn.Sequential()
generator:add( encoder )
generator:add( decoder )
------------------------------------------------------------------
-- discriminator
discriminator = nn.Sequential()
discriminator:add( nn.SpatialConvolution( 3, 32, 4, 4, 2, 2, 1, 1 ) ) -- 3x128x128 -> 32x64x64
-- discriminator.modules[#discriminator.modules].weight:normal(0, 0.02)
-- discriminator.modules[#discriminator.modules].bias:fill(0.1)
discriminator:add( nn.SpatialBatchNormalization( 32 ) )
-- discriminator.modules[#discriminator.modules].weight:normal(1.0, 0.02)
-- discriminator.modules[#discriminator.modules].bias:fill(0)
discriminator:add(nn.LeakyReLU(0.2, true))
discriminator:add( nn.SpatialConvolution( 32, 64, 4, 4, 2, 2, 1, 1 ) ) -- 32x64x64 -> 64x32x32
-- discriminator.modules[#discriminator.modules].weight:normal(0, 0.02)
-- discriminator.modules[#discriminator.modules].bias:fill(0.1)
discriminator:add( nn.SpatialBatchNormalization( 64 ) )
-- discriminator.modules[#discriminator.modules].weight:normal(1.0, 0.02)
-- discriminator.modules[#discriminator.modules].bias:fill(0)
discriminator:add(nn.LeakyReLU(0.2, true))
discriminator:add( nn.SpatialConvolution( 64, 128, 4, 4, 2, 2, 1, 1 ) ) -- 64x32x32 -> 128x16x16
-- discriminator.modules[#discriminator.modules].weight:normal(0, 0.02)
-- discriminator.modules[#discriminator.modules].bias:fill(0.1)
discriminator:add( nn.SpatialBatchNormalization( 128 ) )
-- discriminator.modules[#discriminator.modules].weight:normal(1.0, 0.02)
-- discriminator.modules[#discriminator.modules].bias:fill(0)
discriminator:add(nn.LeakyReLU(0.2, true))
discriminator:add( nn.SpatialConvolution( 128, 256, 4, 4, 2, 2, 1, 1 ) ) -- 128x16x16 -> 256x8x8
-- discriminator.modules[#discriminator.modules].weight:normal(0, 0.02)
-- discriminator.modules[#discriminator.modules].bias:fill(0.1)
discriminator:add( nn.SpatialBatchNormalization( 256 ) )
-- discriminator.modules[#discriminator.modules].weight:normal(1.0, 0.02)
-- discriminator.modules[#discriminator.modules].bias:fill(0)
discriminator:add(nn.LeakyReLU(0.2, true))
discriminator:add( nn.SpatialConvolution( 256, 512, 4, 4, 2, 2, 1, 1 ) ) -- 256x8x8 -> 512x4x4
-- discriminator.modules[#discriminator.modules].weight:normal(0, 0.02)
-- discriminator.modules[#discriminator.modules].bias:fill(0.1)
discriminator:add( nn.View( 512*4*4 ) )
discriminator:add( nn.Linear(512*4*4, 1) )
discriminator:add( nn.Sigmoid() )
--converting to cuda
--classifier:cuda()
generator:cuda()
discriminator:cuda()