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