Module.lua

local _ = require 'moses'

local Module = nn.Module

function Module:sparseParameters()
 return self:parameters()
end

function Module:updateParameters(learningRate)
 -- sparse params can have different learningRate scales per param
 local params, gradParams, scales = self:sparseParameters()
 if params then
 for i,param in pairs(params) do -- pairs for sparse params
 local scale = scales and scales[i] or 1
 param:add(-learningRate*scale, gradParams[i])
 end
 end
end

function Module:zeroGradParameters()
 local _,gradParams = self:sparseParameters()
 if gradParams then
 for i,gradParam in pairs(gradParams) do -- pairs for sparse params
 gradParam:zero()
 end
 end
end

------------------------ clone and type --------------------------------

Module.dpnn_parameters = {'weight', 'bias'}
Module.dpnn_gradParameters = {'gradWeight', 'gradBias'}

-- efficient version of :
-- clone = self:clone()
-- clone:share(self, paramNames, gradParamNames)
-- Note that this method is the very bane of my existence. 
-- I have worked on it too many times...
function Module:sharedClone(shareParams, shareGradParams, stepClone) 
 shareParams = (shareParams == nil) and true or shareParams
 shareGradParams = (shareGradParams == nil) and true or shareGradParams
 
 if stepClone and self.dpnn_stepclone then
 -- this is for AbstractRecurrent modules (in rnn)
 return self
 end
 
 local pointers = {} -- to params/gradParams (dont clone params/gradParams)
 local scdone = {}
 
 -- 1. remove all params/gradParams 
 local function recursiveRemove(obj) -- remove modules
 local moduleTree
 local isTable = type(obj) == 'table' 
 if torch.isTypeOf(obj, 'nn.Module') then
 assert(isTable)
 if stepClone and obj.dpnn_stepclone then
 -- this is for AbstractRecurrent modules (in rnn)
 moduleTree = obj
 obj = nil
 isTable = false
 elseif obj.dpnn_sharedClone then
 -- allow to use a custom sharedClone method on one module
 moduleTree = obj
 obj = nil
 isTable = false
 elseif scdone[torch.pointer(obj)] then
 moduleTree = scdone[torch.pointer(obj)]
 else
 -- remove the params, gradParams. Save for later.
 local params = {}
 
 if shareParams then
 for i,paramName in ipairs(obj.dpnn_parameters) do
 local param = obj[paramName]
 if param then
 params[paramName] = param
 obj[paramName] = nil
 if torch.isTensor(param) and param.storage and param:storage() then
 pointers[torch.pointer(param:storage():data())] = true
 end
 end
 end
 end
 
 if shareGradParams then
 for i,paramName in ipairs(obj.dpnn_gradParameters) do
 local gradParam = obj[paramName]
 if gradParam then
 params[paramName] = gradParam
 obj[paramName] = nil
 if torch.isTensor(gradParam) and gradParam.storage and gradParam:storage() then
 pointers[torch.pointer(gradParam:storage():data())] = true
 end
 end
 end
 end
 
 -- find all obj.attribute tensors that share storage with the shared params
 for paramName, param in pairs(obj) do
 if torch.isTensor(param) and param:storage() then
 if pointers[torch.pointer(param:storage():data())] then
 params[paramName] = param
 obj[paramName] = nil
 end
 end
 end
 
 moduleTree = params
 
 scdone[torch.pointer(obj)] = moduleTree
 
 for k,v in pairs(obj) do
 moduleTree[k], obj[k] = recursiveRemove(v)
 end
 
 end
 elseif isTable then
 if scdone[torch.pointer(obj)] then
 moduleTree = scdone[torch.pointer(obj)]
 else
 assert(not moduleTree)
 moduleTree = {}
 for k,v in pairs(obj) do
 moduleTree[k], obj[k] = recursiveRemove(v)
 end 
 scdone[torch.pointer(obj)] = moduleTree
 end
 
 end
 
 return moduleTree, obj
 end
 
 local moduleTree, original = recursiveRemove(self)
 assert(original)
 
 -- 2. clone everything but parameters, gradients and modules (removed above)
 
 local clone = self:clone()
 
 -- 3. add back to self/clone everything that was removed in step 1
 
 local function recursiveSet(clone, original, moduleTree)
 assert(clone)
 assert(original)
 if scdone[torch.pointer(original)] then
 for k,param in pairs(moduleTree) do
 if torch.isTypeOf(param,'nn.Module') then
 if param.dpnn_sharedClone then
 -- Call the custom sharedClone
 clone[k] = param:dpnn_sharedClone()
 else
 -- AbstractRecurrent instances branch here with stepClone = true
 clone[k] = param
 end
 original[k] = param
 elseif torch.isTensor(param) then
 if param.storage then
 clone[k] = param.new():set(param)
 original[k] = param
 else -- for torch.MultiCudaTensor
 clone[k] = param
 original[k] = param
 end
 elseif type(param) == 'table' then
 recursiveSet(clone[k], original[k], param)
 end
 end 
 scdone[torch.pointer(original)] = nil
 end
 
 end
 
 recursiveSet(clone, self, moduleTree)
 
 return clone
end 

-- we override this method such that hidden modules
-- will be included in the getParameters call.
-- Hidden modules are common for recurrent modules that
-- have internal references to modules that share parameters 
-- with the main modules.
-- These must also be included in the getParameters() call in order 
-- to maintain shared storage for tensors.
function Module:getParameters()

 local con = nn.Container()
 con:add(self)
 
 -- recursive get all modules (modules, sharedclones, etc.)
 local function recursiveGetModules(tbl)
 for k,m in pairs(tbl) do
 if torch.isTypeOf(m, 'nn.Module') then
 if not m.dpnn_getParameters_found then
 con:add(m)
 m.dpnn_getParameters_found = true
 recursiveGetModules(m)
 end
 elseif torch.type(m) == 'table' then
 recursiveGetModules(m)
 end
 end
 end
 
 recursiveGetModules(self)
 
 for i,m in ipairs(con.modules) do
 m.dpnn_getParameters_found = nil
 end

 -- get ALL parameters
 local parameters,gradParameters = con:parameters()
 return Module.flatten(parameters), Module.flatten(gradParameters)
end

----------------- serialization (see nn.Serial) -------------------

Module.dpnn_mediumEmpty = {'output', 'gradInput', 'momGradParams', 'dpnn_input'}
Module.dpnn_lightEmpty = Module.dpnn_gradParameters
-- defaults to heavy serialization
Module.dpnn_serialEmpty = {}

-- sets the serialization behavior of the entire module structure
function Module:serialMode(empty)
 assert(torch.type(empty) == 'table', "Expecting table at arg 1")
 self.dpnn_serialEmpty = empty
 -- set the serial of all encapsulated modules
 local function recursiveSerial(tbl)
 for k,v in pairs(tbl) do
 if torch.isTypeOf(v, 'nn.Module') then
 v:serialMode(empty)
 elseif torch.type(v) == 'table' then
 recursiveSerial(v)
 end
 end
 end
 recursiveSerial(self)
 return self
end

-- serialMode : serialize everything
function Module:heavySerial()
 return self:serialMode({})
end

-- serialMode : serialize everything except dpnn_mediumEmpty attributes
function Module:mediumSerial()
 
 self.dpnn_serialEmpty = self.dpnn_mediumEmpty
 
 -- set the serial of all encapsulated modules
 local function recursiveSerial(tbl)
 for k,v in pairs(tbl) do
 if torch.isTypeOf(v, 'nn.Module') then
 v:mediumSerial()
 elseif torch.type(v) == 'table' then
 recursiveSerial(v)
 end
 end
 end
 recursiveSerial(self)
 return self
end

-- serialMode : serialize everything except dpnn_mediumEmpty and dpnn_lightEmpty attributes
function Module:lightSerial()
 
 self.dpnn_serialEmpty = _.clone(self.dpnn_mediumEmpty)
 for k,v in ipairs(self.dpnn_lightEmpty) do
 table.insert(self.dpnn_serialEmpty, v)
 end
 
 -- set the serial of all encapsulated modules
 local function recursiveSerial(tbl)
 for k,v in pairs(tbl) do
 if torch.isTypeOf(v, 'nn.Module') then
 v:lightSerial()
 elseif torch.type(v) == 'table' then
 recursiveSerial(v)
 end
 end
 end
 recursiveSerial(self)
 
 return self
end

function Module:getSerialState(states)
 states = states or {}
 
 -- dont get the serial state of the same module twice (reuse existing)
 if states[self] then
 return states[self]
 end
 
 -- returns the object structure as tables (i.e. without metatables)
 local function recursiveState(tbl)
 local state = _.map(tbl, 
 function(k,v) 
 if torch.isTypeOf(tbl, 'nn.Module') and _.contains(tbl.dpnn_serialEmpty, k) then 
 -- "empties" module attributes found in empty
 if torch.type(v) == 'table' then
 -- empty table
 return {} 
 elseif torch.isTensor(v) then
 -- empty tensor
 return v.new() 
 else
 -- not table nor tensor? then serialize as is
 return v
 end
 elseif torch.isTypeOf(v, 'nn.Module') then
 -- recursive, yet can be overwritten
 return v:getSerialState(states)
 elseif torch.type(v) == 'table' then
 -- in case it is a table of modules
 if not states[v] then
 states[v] = recursiveState(v)
 end
 return states[v]
 else
 return v
 end
 end
 )
 return state
 end
 
 local state = recursiveState(self)
 
 -- include typename so that module can be reconstructed from the state
 state.dpnn_typename = torch.type(self)
 states[self] = state
 
 return state
end

-- decorates self with nn.Serial
function Module:Serial(tensortype)
 return nn.Serial(self, tensortype)
end

----------------------- for training -----------------------------

-- useful to get the output size
-- I chose this method name because it is less likely to be overriden.
function Module:outside(insize)
 local input
 if torch.type(insize) == 'table' then
 input = torch.randn(table.unpack(insize))
 else
 input = torch.randn(insize)
 end
 local output = self:updateOutput(input)
 return output:size()
end

-- for those interested in implementing the visitor design pattern
function Module:accept(visitor)
 visitor:visit(self)
end

-- Can be used as a regularizer instead of weight decay
-- Assumes that parameters are arranged (output dim x ... x input dim)
function Module:maxParamNorm(maxOutNorm, maxInNorm)
 -- this allows each module to set its own max[Out,In]Norm
 maxOutNorm = self.maxOutNorm or maxOutNorm
 maxInNorm = self.maxInNorm or maxInNorm
 if not (maxOutNorm or maxInNorm) then
 return
 end
 
 if self.modules then
 for i,module in ipairs(self.modules) do
 module:maxParamNorm(maxOutNorm, maxInNorm)
 end
 else
 local params = self:parameters() 
 if not params or gradParams then
 return
 end
 for k,param in pairs(params) do -- pairs for sparse params
 -- By default, only affects non-1D params.
 if param:dim() > 1 then
 if maxOutNorm and maxOutNorm > 0 then
 -- rows feed into output neurons 
 param:renorm(2, 1, maxOutNorm)
 end
 if maxInNorm and maxInNorm > 0 then
 -- cols feed out from input neurons
 param:renorm(2, param:dim(), maxInNorm)
 end
 end
 end
 end
end

-- Similar to maxParamNorm, but norm is global to Module for which 
-- this is called. Unless moduleLocal is true, in which case, the
-- norm constraint is applied to the norm of all parameters in each
-- component (non-container) module.
function Module:gradParamClip(cutoffNorm, moduleLocal)
 -- this allows each module to set its own cutoffNorm
 cutoffNorm = self.cutoffNorm or cutoffNorm
 if cutoffNorm <= 0 then
 return
 end
 if self.moduleLocal ~= nil then
 moduleLocal = self.moduleLocal
 end
 
 local norm = 0
 if moduleLocal and self.modules then
 for i,module in ipairs(self.modules) do
 norm = norm + math.pow(module:gradParamClip(cutoffNorm, moduleLocal), 2)
 end
 norm = math.sqrt(norm)
 else
 local params, gradParams = self:parameters()
 if not (params and gradParams) then
 return norm
 end
 for k,gradParam in pairs(gradParams) do -- pairs for sparse params
 if torch.type(gradParam) == 'torch.CudaTensor' then
 cutorch.withDevice(gradParam:getDevice(), function() -- support multi-device models
 norm = norm + math.pow(gradParam:norm(),2)
 end)
 else
 norm = norm + math.pow(gradParam:norm(),2)
 end
 end
 norm = math.sqrt(norm)
 if norm > cutoffNorm then
 -- rescale gradParams to obtain desired cutoffNorm
 for k,gradParam in pairs(gradParams) do
 if torch.type(gradParam) == 'torch.CudaTensor' then
 cutorch.withDevice(gradParam:getDevice(), function() -- support multi-device models
 gradParam:mul(cutoffNorm/norm)
 end)
 else
 gradParam:mul(cutoffNorm/norm)
 end
 end
 end
 end
 return norm
end

-- Adds weight decay constraint on params with dims > 2 (default).
-- TODO : allow inplace weightDecay (before calling accUpdateGradParameters)
function Module:weightDecay(wdFactor, wdMinDim)
 -- this allows each module to set its own hyper-parameters
 wdFactor = self.wdFactor or wdFactor
 if wdFactor <= 0 then
 return
 end
 wdMinDim = self.wdMinDim or wdMinDim or 2
 
 if self.modules then
 for i,module in ipairs(self.modules) do
 module:weightDecay(wdFactor, wdMinDim)
 end
 else
 local params, gradParams = self:parameters()
 if not (params and gradParams) then
 return
 end
 
 for i,param in pairs(params) do -- pairs for sparse params
 if param:dim() >= wdMinDim then
 gradParams[i]:add(wdFactor, param)
 end
 end
 end
end

function Module:momentumGradParameters()
 if (not self.momGradParams) or _.isEmpty(self.momGradParams) then
 local params, gradParams = self:parameters()
 if not gradParams or _.isEmpty(gradParams) then
 return
 end
 self.momGradParams = {}
 for i,gradParam in pairs(gradParams) do 
 if torch.type(gradParam) == 'torch.CudaTensor' then
 cutorch.withDevice(gradParam:getDevice(), function() -- support multi-device models
 self.momGradParams[i] = gradParam.new():resizeAs(gradParam):copy(gradParam)
 end)
 else
 self.momGradParams[i] = gradParam.new():resizeAs(gradParam):copy(gradParam)
 end
 end
 end
 return self.momGradParams
end

-- uses momentum learning to update gradParams
function Module:updateGradParameters(momFactor, momDamp, momNesterov)
 -- this allows each module to set its own hyper-parameters
 momFactor = self.momFactor or momFactor
 if momFactor <= 0 then
 return
 end
 momDamp = self.momDamp or momDamp or momFactor
 if self.momNesterov ~= nil then
 momNesterov = self.momNesterov
 end
 
 if self.modules then
 for i,module in ipairs(self.modules) do
 module:updateGradParameters(momFactor, momDamp, momNesterov)
 end
 else
 local params, gradParams = self:parameters()
 if (not params) or _.isEmpty(params) then
 return
 end
 local momGradParams = self:momentumGradParameters()
 for i,gradParam in pairs(gradParams) do
 momGradParams[i]:mul(momFactor)
 momGradParams[i]:add(1-momDamp, gradParam)
 end
 
 if momNesterov then
 for i,gradParam in pairs(gradParams) do
 gradParam:add(momFactor, momGradParams[i])
 end
 else
 for i,gradParam in pairs(gradParams) do
 gradParam:copy(momGradParams[i])
 end
 end
 end
end

function Module:checkParameters()
 local params = self:parameters() or {}
 for k,param in pairs(params) do
 if _.isNaN(param:sum()) then
 error("NaN Error for param at index" ..k)
 end
 end
end

function Module:dontBackward()
 self.backward = function() end
 self.updateGradInput = function() end
 self.accGradParameters = function() end
 self.accUpdateGradParameters = function() end
 return self
end

function Module:contiguousInput(input, backward)
 if backward then
 return self.dpnn_cinput or input
 end
 if not input:isContiguous() then
 self.dpnn_cinput = self.dpnn_cinput or input.new()
 self.dpnn_cinput:resizeAs(input):copy(input)
 input = self.dpnn_cinput
 end
 return input
end

function Module:toBatch(tensor, nDim, batchDim)
 local batchDim = batchDim or 1
 if tensor:dim() == nDim then
 self.dpnn_online = true
 local size = tensor:size():totable()
 table.insert(size, batchDim, 1)
 tensor = tensor:view(table.unpack(size))
 else
 self.dpnn_online = false
 end
 return tensor
end

function Module:fromBatch(tensor, batchDim)
 if self.dpnn_online then
 local size = tensor:size():totable()
 assert(table.remove(size, batchDim) == 1)
 tensor = tensor:view(table.unpack(size))
 end
 return tensor
end

function Module:extrapolateType()
 local params = module:parameters()
 if params then
 -- extrapolate the tensor type of the module
 local types = {}
 for i, param in ipairs(params) do
 local tensorType = torch.type(param)
 types[tensorType] = (types[tensorType] or 0) + 1
 end
 local maxCount = 0
 local maxType
 for tensorType, count in pairs(types) do
 if count > maxCount then
 maxtype = tensorType
 maxCount = count
 end
 end
 return maxType
 end
 return nil --unknown otherwise
end

function Module:profile()
 if self.modules then
 for i, module in ipairs(self.modules) do
 module:profile()
 end
 end
 self.dpnn_profile = true
end

function Module:reinforce(reward)
 if self.modules then
 for i, module in ipairs(self.modules) do
 module:reinforce(reward)
 end
 end
end