maptable.lua

require 'nn'
require 'nngraph'

local MapTable, parent = torch.class('nn.maptable', 'nn.Container')

function MapTable:__init(module, shared,num_inputs)
   parent.__init(self)
   self.shared = (shared == nil) and true or shared
   self.sharedparams = {'weight', 'bias', 'gradWeight', 'gradBias'}
   self.output = {}
   self.gradInput = {}
   self:add(module)
   self:_extend(num_inputs)
end

function MapTable:_extend(n)
   self.sharedparams = self.sharedparams or {'weight', 'bias', 'gradWeight', 'gradBias'}
   self.modules[1] = self.module
   --assert(n==3,"provide 3 inputs")
   p,g=self.module:getParameters()
   print("size of original",p:size(1))
   for num = 2, n do
      if not self.modules[i] then
         if self.shared then
           self.modules[num] = self.module:clone()
           for i=1,#self.module.forwardnodes do
             if(self.module.forwardnodes[i].data.module~=nil and self.module.forwardnodes[i].data.module.weight~=nil) then
                self.module.forwardnodes[i].data.module:share(self.modules[num].forwardnodes[i].data.module, 'weight','bias','gradWeight','gradBias');
                -- print(self.module.forwardnodes[i].data.annotations.name)
             end 
             if(self.module.forwardnodes[i].data.annotations.name=='attention_module') then
              atten1=self.module.forwardnodes[i].data.module.lstm
              atten2=self.modules[num].forwardnodes[i].data.module.lstm
              for j=1,#atten1.forwardnodes do
                if(atten1.forwardnodes[j].data.module~=nil and atten1.forwardnodes[j].data.module.weight~=nil) then
                  atten1.forwardnodes[j].data.module:share(atten2.forwardnodes[j].data.module, 'weight','bias','gradWeight','gradBias');
                  -- print(atten1.forwardnodes[j].data.annotations.name)
                end
              end 
             end  
            end 
         else
           self.modules[num] = self.module:clone()
         end
      end
      p,g=self.modules[num]:getParameters()
      print(p:size(1),num)
   end
end

function MapTable:resize(n)
   -- self:_extend(n)
   for i = n + 1, #self.modules do
      -- It's not clear why this clearState call is necessary, but it fixes
      -- https://github.com/torch/nn/issues/1141 .
      self.modules[i]:clearState()
      self.modules[i] = nil
   end
end

function MapTable:add(module)
   assert(not self.module, 'Single module required')
   self.module = module
   self.modules[1] = self.module
   return self
end

function MapTable:updateOutput(input)
   self.output = {}
   -- self:_extend(#input)
   for i = 1, #input do
      self.output[i] = self:rethrowErrors(self.modules[i], i, 'updateOutput', input[i])
   end
   return self.output
end

function MapTable:updateGradInput(input, gradOutput)
   self.gradInput = {}
   -- self:_extend(#input)
   for i = 1, #input do
      self.gradInput[i] = self:rethrowErrors(self.modules[i], i, 'updateGradInput', input[i], gradOutput[i])
   end
   return self.gradInput
end

function MapTable:accGradParameters(input, gradOutput, scale)
   scale = scale or 1
   -- self:_extend(#input)
   for i = 1, #input do
      self:rethrowErrors(self.modules[i], i, 'accGradParameters', input[i], gradOutput[i], scale)
   end
end

function MapTable:accUpdateGradParameters(input, gradOutput, lr)
   lr = lr or 1
   -- self:_extend(#input)
   for i = 1, #input do
      self:rethrowErrors(self.modules[i], i, 'accUpdateGradParameters', input[i], gradOutput[i], lr)
   end
end

function MapTable:zeroGradParameters()
    if self.module then
        if self.shared then
          self.module:zeroGradParameters()
        else
          parent.zeroGradParameters(self)
        end
    end
end

function MapTable:updateParameters(learningRate)
    if self.module then
        if self.shared then
          self.module:updateParameters(learningRate)
        else
          parent.updateParameters(self, learningRate)
        end
    end
end

function MapTable:clearState()
   for i = 2, #self.modules do
      -- It's not clear why this clearState call is necessary, but it fixes
      -- https://github.com/torch/nn/issues/1141 .
      self.modules[i]:clearState()
      self.modules[i] = nil
   end
   parent.clearState(self)
end

function MapTable:__tostring__()
   local tab = '  '
   local line = '\n'
   local extlast = '      '
   local str = torch.type(self)
   if self.module then
      str = str .. ' {' .. line .. tab
      str = str .. tostring(self.module):gsub(line, line .. tab .. extlast) .. line .. '}'
   else
      str = str .. ' { }'
   end
   return str
end