This repository has been archived by the owner on Jul 12, 2018. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 12
/
Copy pathmodel.lua
93 lines (78 loc) · 2.84 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
require('nn')
require('image')
torch.setdefaulttensortype('torch.FloatTensor')
local classNumber = 21
local allBoxes = { {2,2}, {3,3}, {4,4}, {5,5},
{2,4}, {4,2}, {3,6}, {6,3},
{2,6}, {6,2},
{2,3}, {3,2}, {4,6}, {6,4},
{2,5}, {5,2}, {3,4}, {4,3},
{3,5}, {5,3}, {4,5}, {5,4} }
local fixedCNN = torch.load('fixedCNN.t7');
fixedCNN:evaluate()
local featureCNN = nn.Sequential()
-- input is 56x56
featureCNN:add( nn.SpatialConvolution(128, 256, 3, 3, 1, 1, 1, 1) )
featureCNN:add( nn.ReLU(true) )
featureCNN:add( nn.SpatialConvolution(256, 256, 3, 3, 1, 1, 1, 1) )
featureCNN:add( nn.ReLU(true) )
featureCNN:add( nn.SpatialConvolution(256, 256, 3, 3, 1, 1, 1, 1) )
featureCNN:add( nn.ReLU(true) )
featureCNN:add( nn.SpatialMaxPooling(2, 2, 2, 2, 0, 0):ceil() )
-- input is 28x28
featureCNN:add( nn.SpatialConvolution(256, 512, 3, 3, 1, 1, 1, 1) )
featureCNN:add( nn.ReLU(true) )
featureCNN:add( nn.SpatialConvolution(512, 512, 3, 3, 1, 1, 1, 1) )
featureCNN:add( nn.ReLU(true) )
featureCNN:add( nn.SpatialConvolution(512, 512, 3, 3, 1, 1, 1, 1) )
featureCNN:add( nn.ReLU(true) )
featureCNN:add( nn.SpatialMaxPooling(2, 2, 2, 2, 0, 0):ceil() )
-- input is 14x14
featureCNN:add( nn.SpatialConvolution(512, 512, 3, 3, 1, 1, 1, 1) )
featureCNN:add( nn.ReLU(true) )
featureCNN:add( nn.SpatialConvolution(512, 512, 3, 3, 1, 1, 1, 1) )
featureCNN:add( nn.ReLU(true) )
--[[
featureCNN:add( nn.SpatialConvolution(512, 512, 3, 3, 1, 1, 1, 1) )
featureCNN:add( nn.ReLU(true) )
featureCNN:add( nn.SpatialMaxPooling(2, 2, 2, 2, 0, 0):ceil() )
--]]
featureCNN:add(nn.SpatialConvolution(512, 1024, 1, 1, 1, 1, 0, 0))
featureCNN:add(nn.LeakyReLU(0.1))
local lossLayers = {}
local mbox = nn.ConcatTable()
for i = 1, #allBoxes do
local boxConf = nn.Sequential()
boxConf:add(nn.Dropout(0.5))
boxConf:add(nn.SpatialConvolution(1024, classNumber, allBoxes[i][1], allBoxes[i][2], 1, 1, 0, 0))
boxConf:add(nn.SpatialLogSoftMax())
mbox:add(boxConf)
table.insert(lossLayers, nn.SpatialClassNLLCriterion())
local boxLoc = nn.Sequential()
boxLoc:add(nn.SpatialConvolution(1024, 4, allBoxes[i][1], allBoxes[i][2], 1, 1, 0, 0))
mbox:add(boxLoc)
table.insert(lossLayers, nn.MSECriterion())
end
featureCNN:add(mbox)
--[[
fixedCNN:cuda()
featureCNN:cuda()
local x = torch.rand(4, 3,256,256):cuda()
local y = featureCNN:forward( fixedCNN:forward(x) )
print(y)
--]]
local getSize = function(imageWidth, imageHeight)
local targetWidth = math.floor(imageWidth/16)
local targetHeight = math.floor(imageHeight/16)
return {targetWidth, targetHeight};
end
local model = {}
model.fixedCNN = fixedCNN
model.featureCNN = featureCNN
model.lossLayers = lossLayers
local info = {}
info.classNumber = classNumber
info.boxes = allBoxes
info.getSize = getSize
model.info = info
return model