-
Notifications
You must be signed in to change notification settings - Fork 8
/
delugenet.lua
181 lines (162 loc) · 5.27 KB
/
delugenet.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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
--
-- DelugeNet
--
local nn = require 'nn'
require 'cunn'
require 'models/CrossLayerDepthwiseConvolution'
local Convolution = cudnn.SpatialConvolution
local Avg = cudnn.SpatialAveragePooling
local ReLU = cudnn.ReLU
local Max = cudnn.SpatialMaxPooling
local SBatchNorm = cudnn.SpatialBatchNormalization
local function createModel(opt)
local depth = opt.depth
local iChannels, nLayers = nil, 1
-- Bottleneck composite layer
local function bottleneck(n, stride, type)
local nInputPlane = iChannels
iChannels = n * 2
local block = nn.Sequential()
local s = nn.Sequential()
if type == 'both_preact' or type ~= 'no_preact' then
if nLayers == 1 then
s:add(SBatchNorm(nInputPlane))
elseif nLayers > 1 then
s:add(nn.CrossLayerDepthwiseConvolution(nInputPlane, nLayers))
s:add(SBatchNorm(nInputPlane))
-- block transition
if type ~= 'no_preact' and type ~= nil then
nLayers = 1
s:add(Convolution(nInputPlane,n*2,3,3,stride,stride,1,1))
return block:add(s)
end
end
s:add(ReLU(true))
elseif type == 'no_preact' then
s:add(Convolution(nInputPlane,n*2,1,1,1,1,0,0))
return block:add(s)
end
s:add(Convolution(nInputPlane,n,1,1,1,1,0,0))
s:add(SBatchNorm(n))
s:add(ReLU(true))
s:add(Convolution(n,n,3,3,stride,stride,1,1))
s:add(SBatchNorm(n))
s:add(ReLU(true))
s:add(Convolution(n,n*2,1,1,1,1,0,0))
if type == nil then
nLayers = nLayers + 1
return block
:add(nn.ConcatTable()
:add(s)
:add(nn.Identity()))
:add(nn.FlattenTable())
else
nLayers = 1
return block:add(s)
end
end
-- Network block
local function block(features, count, stride, type)
local s = nn.Sequential()
if count < 1 then
return s
end
s:add(bottleneck(features, stride,
type == 'first' and 'no_preact' or 'both_preact'))
for i=2,count do
s:add(bottleneck(features, 1, nil))
end
return s
end
local model = nn.Sequential()
if opt.dataset == 'imagenet' then
-- Configurations:
-- compositeLayer counts
local cfg = {
[92] = {7+1, 7+1, 8+1, 8+1},
[104] = {7+1, 8+1, 9+1, 10+1},
[122] = {7+1, 9+1, 11+1, 13+1},
}
assert(cfg[depth], 'Invalid depth: ' .. tostring(depth))
local def = table.unpack(cfg[depth])
iChannels = 64
print(' | DelugeNet-' .. depth .. ' ImageNet')
-- The ImageNet model
model:add(Convolution(3,64,7,7,2,2,3,3))
model:add(SBatchNorm(64))
model:add(ReLU(true))
model:add(Max(3,3,2,2,1,1))
model:add(block(64, def[1], 1, 'first'))
model:add(block(128, def[2], 2))
model:add(block(256, def[3], 2))
model:add(block(512, def[4], 2))
model:add(nn.CrossLayerDepthwiseConvolution(iChannels, nLayers))
model:add(SBatchNorm(iChannels))
model:add(ReLU(true))
model:add(Avg(7, 7, 1, 1))
model:add(nn.View(iChannels):setNumInputDims(3))
model:add(nn.Linear(iChannels, 1000))
elseif opt.dataset == 'cifar10' or opt.dataset == 'cifar100' then
-- Configurations:
-- # compositeLayer in first block, width
local cfg = {
[146.1] = {8,1},
[146.2] = {8,1.75},
[218] = {12,1}
}
assert(cfg[depth], 'Invalid depth: ' .. tostring(depth))
local n, width = table.unpack(cfg[depth])
iChannels = 32*width
print(' | DelugeNet-' .. depth .. ' ' .. (opt.dataset == 'cifar10' and 'CIFAR-10' or 'CIFAR-100'))
-- The CIFAR model
model:add(Convolution(3,32*width,3,3,1,1,1,1))
model:add(block(32*width, n, 1, nil))
model:add(block(64*width, n*2+1, 2, nil))
model:add(block(128*width, n*3+1, 2, nil))
model:add(nn.CrossLayerDepthwiseConvolution(iChannels, nLayers))
model:add(SBatchNorm(iChannels))
model:add(ReLU(true))
model:add(Avg(8, 8, 1, 1))
model:add(nn.View(iChannels):setNumInputDims(3))
model:add(nn.Linear(iChannels, opt.dataset == 'cifar10' and 10 or 100))
else
error('invalid dataset: ' .. opt.dataset)
end
local function ConvInit(name)
for k,v in pairs(model:findModules(name)) do
local n = v.kW*v.kH*v.nOutputPlane
v.weight:normal(0,math.sqrt(2/n))
if cudnn.version >= 4000 then
v.bias = nil
v.gradBias = nil
else
v.bias:zero()
end
end
end
local function BNInit(name)
for k,v in pairs(model:findModules(name)) do
if v.weight then
v.weight:fill(1)
v.bias:zero()
end
end
end
ConvInit('cudnn.SpatialConvolution')
ConvInit('nn.SpatialConvolution')
BNInit('fbnn.SpatialBatchNormalization')
BNInit('cudnn.SpatialBatchNormalization')
BNInit('nn.SpatialBatchNormalization')
for k,v in pairs(model:findModules('nn.Linear')) do
v.bias:zero()
end
model:cuda()
if opt.cudnn == 'deterministic' then
model:apply(function(m)
if m.setMode then m:setMode(1,1,1) end
end)
end
model:get(1).gradInput = nil
return model
end
return createModel