Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
177 lines (155 sloc) 5.74 KB
--[[
Copyright (c) 2016 Michael Wilber
This software is provided 'as-is', without any express or implied
warranty. In no event will the authors be held liable for any damages
arising from the use of this software.
Permission is granted to anyone to use this software for any purpose,
including commercial applications, and to alter it and redistribute it
freely, subject to the following restrictions:
1. The origin of this software must not be misrepresented; you must not
claim that you wrote the original software. If you use this software
in a product, an acknowledgement in the product documentation would be
appreciated but is not required.
2. Altered source versions must be plainly marked as such, and must not be
misrepresented as being the original software.
3. This notice may not be removed or altered from any source distribution.
--]]
require 'nn'
require 'nngraph'
require 'cudnn'
require 'cunn'
local nninit = require 'nninit'
function addResidualLayer2(input, nChannels, nOutChannels, stride)
--[[
Residual layers! Implements option (A) from Section 3.3. The input
is passed through two 3x3 convolution filters. In parallel, if the
number of input and output channels differ or if the stride is not
1, then the input is downsampled or zero-padded to have the correct
size and number of channels. Finally, the two versions of the input
are added together.
Input
|
,-------+-----.
Downsampling 3x3 convolution+dimensionality reduction
| |
v v
Zero-padding 3x3 convolution
| |
`-----( Add )---'
|
Output
--]]
nOutChannels = nOutChannels or nChannels
stride = stride or 1
-- Path 1: Convolution
-- The first layer does the downsampling and the striding
local net = cudnn.SpatialConvolution(nChannels, nOutChannels,
3,3, stride,stride, 1,1)
:init('weight', nninit.kaiming, {gain = 'relu'})
:init('bias', nninit.constant, 0)(input)
net = cudnn.SpatialBatchNormalization(nOutChannels)
:init('weight', nninit.normal, 1.0, 0.002)
:init('bias', nninit.constant, 0)(net)
net = cudnn.ReLU(true)(net)
net = cudnn.SpatialConvolution(nOutChannels, nOutChannels,
3,3, 1,1, 1,1)
:init('weight', nninit.kaiming, {gain = 'relu'})
:init('bias', nninit.constant, 0)(net)
-- Should we put Batch Normalization here? I think not, because
-- BN would force the output to have unit variance, which breaks the residual
-- property of the network.
-- What about ReLU here? I think maybe not for the same reason. Figure 2
-- implies that they don't use it here
-- Path 2: Identity / skip connection
local skip = input
if stride > 1 then
-- optional downsampling
skip = nn.SpatialAveragePooling(1, 1, stride,stride)(skip)
end
if nOutChannels > nChannels then
-- optional padding
skip = nn.Padding(1, (nOutChannels - nChannels), 3)(skip)
elseif nOutChannels < nChannels then
-- optional narrow, ugh.
skip = nn.Narrow(2, 1, nOutChannels)(skip)
-- NOTE this BREAKS with non-batch inputs!!
end
-- Add them together
net = cudnn.SpatialBatchNormalization(nOutChannels)(net)
net = nn.CAddTable(){net, skip}
net = cudnn.ReLU(true)(net)
-- ^ don't put a ReLU here! see http://gitxiv.com/comments/7rffyqcPLirEEsmpX
return net
end
--[[
function addResidualLayer3(input, inChannels, hiddenChannels, outChannels)
-- Downsampling and convolution path
local net = cudnn.SpatialConvolution(inChannels, hiddenChannels,
1,1)(input)
net = cudnn.ReLU(true)(net)
net = cudnn.SpatialConvolution(hiddenChannels, hiddenChannels,
3,3, 1,1, 1,1)(net)
--net = nn.Narrow
net = cudnn.ReLU(true)(net)
net = cudnn.SpatialConvolution(hiddenChannels, outChannels,
1,1)(net)
-- Add them together
--return net
return nn.CAddTable(){net, input}
end
--]]
--[[
-- Useful for memory debugging
function countElts(modules)
local sum_elts = 0
for k,v in pairs(modules) do
if torch.isTensor(v) then
sum_elts = sum_elts + v:numel()
elseif torch.type(v) == 'table' then
sum_elts = sum_elts + countElts(v)
end
end
return sum_elts
end
function inspectMemory(net)
local total_count = 0
for i,module in ipairs(net.modules) do
print(i..": "..tostring(module))
local count_this_module = countElts(module)
print(count_this_module)
total_count = total_count + count_this_module
end
print("Total:",total_count)
print(" ",total_count*8/1024./1024., " MB")
end
function accumMemoryByFieldName(module, accum)
for k,v in pairs(module) do
if torch.isTensor(v) then
accum[k] = (accum[k] or 0) + (v:numel() * 8./1024./1024.)
end
end
end
--]]
--[[
-- Testing
input = nn.Identity()()
output = addResidualLayer2(input, 3, 6, 2)
net = nn.gModule({input},{output})
i = torch.randn(1, 3, 5,5):fill(1):cuda()
net:cuda()
net.modules[2].bias:fill(0)
net.modules[4].bias:fill(0)
net.modules[2].weight:fill(0)
net.modules[4].weight:fill(0)
--]]
-- -- Testing memory usage
-- i = torch.randn(5, 256, 224,224)
-- o = net:forward(i)
-- net:backward(i, o)
--
-- inspectMemory(net)
-- mem_usage = {}
-- for i,module in ipairs(net.modules) do
-- accumMemoryByFieldName(module, mem_usage)
-- end
-- print(mem_usage)