dagnn.lua

   1
   2 --[[
   3
   4    Copyright (c) 2016 Idiap Research Institute, http://www.idiap.ch/
   5    Written by Francois Fleuret <francois.fleuret@idiap.ch>
   6
   7    This file is free software: you can redistribute it and/or modify
   8    it under the terms of the GNU General Public License version 3 as
   9    published by the Free Software Foundation.
  10
  11    It is distributed in the hope that it will be useful, but WITHOUT
  12    ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
  13    or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
  14    License for more details.
  15
  16    You should have received a copy of the GNU General Public License
  17    along with this file.  If not, see <http://www.gnu.org/licenses/>.
  18
  19 ]]--
  20
  21 require 'torch'
  22 require 'nn'
  23
  24 local DAG, parent = torch.class('nn.DAG', 'nn.Container')
  25
  26 function DAG:__init()
  27    parent.__init(self)
  28    -- Nodes are indexed by the module they contain
  29    self.node = { }
  30 end
  31
  32 function DAG:createNode(nnm)
  33    if not self.node[nnm] then
  34       self:add(nnm) -- Add it to the object as a Container
  35       self.node[nnm] = {}
  36       self.node[nnm].succ = {}
  37       self.node[nnm].pred = {}
  38    end
  39 end
  40
  41 function DAG:addEdge(nnma, nnmb)
  42    self.sorted = nil
  43    self:createNode(nnma)
  44    self:createNode(nnmb)
  45    table.insert(self.node[nnmb].pred, nnma)
  46    table.insert(self.node[nnma].succ, nnmb)
  47 end
  48
  49 -- Apply f on t recursively; use the corresponding element from args
  50 -- (i.e. same keys) as second parameter to f when available; return
  51 -- the results from f, organized in a similarly nested table.
  52 function DAG:nestedApply(f, t, args)
  53    if torch.type(t) == 'table' then
  54       local result = {}
  55       for k, s in pairs(t) do
  56          result[k] = self:nestedApply(f, s, args and args[k])
  57       end
  58       return result
  59    else
  60       return f(t, args)
  61    end
  62 end
  63
  64 function DAG:setInput(i)
  65    self.sorted = nil
  66    self.inputModules = i
  67    self:nestedApply(
  68       function(nnm)
  69          if #self.node[nnm].succ == 0 then
  70             error('Input modules must have outgoing  edges.')
  71          end
  72          if #self.node[nnm].pred > 0 then
  73             error('Input modules cannog have incoming edges.')
  74          end
  75       end,
  76       self.inputModules
  77    )
  78 end
  79
  80 function DAG:setOutput(o)
  81    self.sorted = nil
  82    self.outputModules = o
  83    self:nestedApply(
  84       function(nnm)
  85          if #self.node[nnm].pred == 0 then
  86             error('Output module must have incoming edges.')
  87          end
  88          if #self.node[nnm].succ > 0 then
  89             error('Output module cannot have outgoing edges.')
  90          end
  91       end,
  92       self.outputModules
  93    )
  94 end
  95
  96 function DAG:putInOrder()
  97    if self.sorted then
  98       return
  99    end
 100
 101    -- First, we sort the nodes according to the DAG order
 102
 103    local distance = {}
 104
 105    self:nestedApply(function(m) distance[m] = 1 end, self.inputModules)
 106
 107    local nc
 108
 109    repeat
 110       nc = 0
 111       for nnma, node in pairs(self.node) do
 112          for _, nnmb in pairs(node.succ) do
 113             if distance[nnma] and (not distance[nnmb] or distance[nnmb] < distance[nnma] + 1) then
 114                distance[nnmb] = distance[nnma] + 1
 115                nc = nc + 1
 116             end
 117          end
 118       end
 119    until nc == 0
 120
 121    self.sorted = { }
 122    for m, d in pairs(distance) do
 123       table.insert(self.sorted, { distance = d, nnm = m })
 124    end
 125
 126    table.sort(self.sorted, function(a, b) return a.distance < b.distance end)
 127
 128    for i, a in ipairs(self.sorted) do self.sorted[i] = a.nnm end
 129 end
 130
 131 function DAG:print()
 132    self:putInOrder()
 133
 134    for i, d in ipairs(self.sorted) do
 135       print('#' .. i .. ' -> ' .. torch.type(d))
 136    end
 137 end
 138
 139 function DAG:updateOutput(input)
 140    self:putInOrder()
 141
 142    self:nestedApply(
 143       function(nnm, i)
 144          self.node[nnm].input = i
 145          nnm:updateOutput(i)
 146       end,
 147       self.inputModules,
 148       input
 149    )
 150
 151    for _, nnm in ipairs(self.sorted) do
 152       local node = self.node[nnm]
 153       if #node.pred > 0 then
 154          local i
 155          if #node.pred == 1 then
 156             i = node.pred[1].output
 157          elseif #node.pred > 1 then
 158             i = {}
 159             for k = 1, #node.pred do
 160                i[k] = node.pred[k].output
 161             end
 162          end
 163          node.input = i
 164          nnm:updateOutput(i)
 165       end
 166    end
 167
 168    self.output = self:nestedApply(
 169       function(m) return m.output end,
 170       self.outputModules
 171    )
 172
 173    return self.output
 174 end
 175
 176 function DAG:computeGradInput(gradInputSucc)
 177    local gi
 178    if #gradInputSucc == 1 then
 179       gi = gradInputSucc[1] -- we avoid a clone()
 180    elseif #gradInputSucc > 1 then
 181       for k = 1, #gradInputSucc do
 182          if gi then
 183             gi:add(gradInputSucc[k])
 184          else
 185             gi = gradInputSucc[k]:clone()
 186          end
 187       end
 188    end
 189    return gi
 190 end
 191
 192 function DAG:updateGradInput(input, gradOutput)
 193    self:putInOrder()
 194
 195    self:nestedApply(
 196       function(nnm, go) nnm:updateGradInput(self.node[nnm].input, go) end,
 197       self.outputModules, gradOutput
 198    )
 199
 200    self:nestedApply(
 201       function(nnm, i) self.node[nnm].input = i end,
 202       self.inputModules, input
 203    )
 204
 205    for _, node in pairs(self.node) do
 206       node.gradInputSucc = {}
 207    end
 208
 209    for k = #self.sorted, 1, -1 do
 210       local nnm = self.sorted[k]
 211       local node = self.node[nnm]
 212       local pred, gradInputSucc = node.pred, node.gradInputSucc
 213
 214       if #gradInputSucc > 0 then
 215          nnm:updateGradInput(node.input, self:computeGradInput(gradInputSucc))
 216       end
 217
 218       -- We fill the gradInputSucc of our predecessors
 219       if #pred == 1 then
 220          table.insert(self.node[pred[1]].gradInputSucc, nnm.gradInput)
 221       elseif #pred > 1 then
 222          if not torch.type(nnm.gradInput) == 'table' then
 223             error('Should have a table gradInput since it has multiple predecessors')
 224          end
 225          for n = 1, #pred do
 226             table.insert(self.node[node.pred[n]].gradInputSucc, nnm.gradInput[n])
 227          end
 228       end
 229    end
 230
 231    self.gradInput = self:nestedApply(function(m) return m.gradInput end, self.inputModules)
 232
 233    return self.gradInput
 234 end
 235
 236 function DAG:accGradParameters(input, gradOutput, scale)
 237    scale = scale or 1
 238
 239    self:putInOrder()
 240
 241    self:nestedApply(
 242       function(nnm, go) nnm:updateGradInput(self.node[nnm].input, go) end,
 243       self.outputModules, gradOutput
 244    )
 245
 246    self:nestedApply(
 247       function(nnm, i) self.node[nnm].input = i end,
 248       self.inputModules, input
 249    )
 250
 251    for k = #self.sorted, 1, -1 do
 252       local nnm = self.sorted[k]
 253       local node = self.node[nnm]
 254       nnm:accGradParameters(node.input, self:computeGradInput(node.gradInputSucc), scale)
 255    end
 256 end