Brian Broll
58 linhas
1.8 KiB
Lua
58 linhas
1.8 KiB
Lua
-- thanks to https://github.com/soumith/imagenet-multiGPU.torch for this example
|
|
require 'nn'
|
|
local nGPU = 4
|
|
local nClasses = 10
|
|
local modelType = 'A' -- on a titan black, B/D/E run out of memory even for batch-size 32
|
|
|
|
-- Create tables describing VGG configurations A, B, D, E
|
|
local cfg = {}
|
|
if modelType == 'A' then
|
|
cfg = {64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'}
|
|
elseif modelType == 'B' then
|
|
cfg = {64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'}
|
|
elseif modelType == 'D' then
|
|
cfg = {64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'}
|
|
elseif modelType == 'E' then
|
|
cfg = {64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'}
|
|
else
|
|
error('Unknown model type: ' .. modelType .. ' | Please specify a modelType A or B or D or E')
|
|
end
|
|
|
|
local features = nn.Sequential()
|
|
do
|
|
local iChannels = 3;
|
|
for k,v in ipairs(cfg) do
|
|
if v == 'M' then
|
|
features:add(nn.SpatialMaxPooling(2,2,2,2))
|
|
else
|
|
local oChannels = v;
|
|
local conv3 = nn.SpatialConvolution(iChannels,oChannels,3,3,1,1,1,1);
|
|
features:add(conv3)
|
|
features:add(nn.ReLU(true))
|
|
iChannels = oChannels;
|
|
end
|
|
end
|
|
end
|
|
|
|
-- features:cuda()
|
|
-- features = makeDataParallel(features, nGPU) -- defined in util.lua
|
|
|
|
local classifier = nn.Sequential()
|
|
classifier:add(nn.View(512*7*7))
|
|
classifier:add(nn.Linear(512*7*7, 4096))
|
|
classifier:add(nn.Threshold(0, 1e-6))
|
|
classifier:add(nn.Dropout(0.5))
|
|
classifier:add(nn.Linear(4096, 4096))
|
|
classifier:add(nn.Threshold(0, 1e-6))
|
|
classifier:add(nn.Dropout(0.5))
|
|
classifier:add(nn.Linear(4096, nClasses))
|
|
classifier:add(nn.LogSoftMax())
|
|
-- classifier:cuda()
|
|
|
|
local model = nn.Sequential()
|
|
model:add(features):add(classifier)
|
|
model.imageSize = 256
|
|
model.imageCrop = 224
|
|
|
|
return model
|