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 -- Apply f on t recursively; use the corresponding elements from args
  33 -- (i.e. same keys) as second parameter to f when available; return
  34 -- the results from f, organized in a similarly nested table.
  35 function DAG:nestedApply(f, t, args)
  36    if torch.type(t) == 'table' then
  37       local result = {}
  38       for k, s in pairs(t) do
  39          result[k] = self:nestedApply(f, s, args and args[k])
  40       end
  41       return result
  42    else
  43       return f(t, args)
  44    end
  45 end
  46
  47 function DAG:createNode(nnm)
  48    if not self.node[nnm] then
  49       self:add(nnm) -- Add it to the object as a Container
  50       local node = {}
  51       node.succ = {}
  52       node.pred = {}
  53       node.index = #self.modules
  54       self.node[nnm] = node
  55    end
  56 end
  57
  58 function DAG:putInOrder()
  59    if self.sorted then
  60       return
  61    end
  62
  63    local distance = {}
  64    self:nestedApply(function(m) distance[m] = 1 end, self.inputModules)
  65
  66    local nc
  67    repeat
  68       nc = 0
  69       for nnma, node in pairs(self.node) do
  70          for _, nnmb in pairs(node.succ) do
  71             if distance[nnma] and (not distance[nnmb] or distance[nnmb] < distance[nnma] + 1) then
  72                distance[nnmb] = distance[nnma] + 1
  73                nc = nc + 1
  74             end
  75          end
  76       end
  77    until nc == 0
  78
  79    self.sorted = { }
  80    for m, d in pairs(distance) do
  81       table.insert(self.sorted, { distance = d, nnm = m })
  82    end
  83
  84    table.sort(self.sorted, function(a, b) return a.distance < b.distance end)
  85
  86    for i, a in ipairs(self.sorted) do self.sorted[i] = a.nnm end
  87 end
  88
  89 function DAG:computeGradOutput(gradInputSucc)
  90    local gi
  91    if #gradInputSucc == 1 then
  92       gi = gradInputSucc[1] -- we avoid a clone()
  93    elseif #gradInputSucc > 1 then
  94       for k = 1, #gradInputSucc do
  95          if gi then
  96             gi:add(gradInputSucc[k])
  97          else
  98             gi = gradInputSucc[k]:clone()
  99          end
 100       end
 101    end
 102    return gi
 103 end
 104
 105 ----------------------------------------------------------------------
 106
 107 -- Connect a sequence of modules
 108 function DAG:connect(...)
 109    self.sorted = nil
 110    local prev
 111    for _, nnm in pairs({...}) do
 112       self:createNode(nnm)
 113       if prev then
 114          table.insert(self.node[nnm].pred, prev)
 115          table.insert(self.node[prev].succ, nnm)
 116       end
 117       prev = nnm
 118    end
 119 end
 120
 121 function DAG:setInput(i)
 122    self.sorted = nil
 123    self.inputModules = i
 124    self:nestedApply(
 125       function(nnm)
 126          if #self.node[nnm].succ == 0 then
 127             error('Input modules must have outgoing  edges.')
 128          end
 129          if #self.node[nnm].pred > 0 then
 130             error('Input modules cannog have incoming edges.')
 131          end
 132       end,
 133       self.inputModules
 134    )
 135 end
 136
 137 function DAG:setOutput(o)
 138    self.sorted = nil
 139    self.outputModules = o
 140    self:nestedApply(
 141       function(nnm)
 142          if #self.node[nnm].pred == 0 then
 143             error('Output module must have incoming edges.')
 144          end
 145          if #self.node[nnm].succ > 0 then
 146             error('Output module cannot have outgoing edges.')
 147          end
 148       end,
 149       self.outputModules
 150    )
 151 end
 152
 153 function DAG:print()
 154    self:putInOrder()
 155
 156    for i, d in ipairs(self.sorted) do
 157       print('#' .. i .. ' -> ' .. torch.type(d))
 158    end
 159 end
 160
 161 ----------------------------------------------------------------------
 162
 163 function DAG:saveDot(filename)
 164    local file = (filename and io.open(filename, 'w')) or io.stdout
 165
 166    file:write('digraph {\n')
 167
 168    file:write('\n')
 169
 170    for nnmb, node in pairs(self.node) do
 171       file:write(
 172          '  '
 173             .. node.index
 174             .. ' [shape=box,label=\"' .. torch.type(nnmb) .. '\"]'
 175             .. '\n'
 176       )
 177
 178       for i, nnma in pairs(node.pred) do
 179          local decoration = ''
 180          if #node.pred > 1 then
 181             -- decoration = ' [headlabel=\"' .. i .. '\"]'
 182             decoration = ' [label=\"' .. i .. '\"]'
 183          end
 184          file:write(
 185             '  '
 186                .. self.node[nnma].index
 187                .. ' -> '
 188                .. self.node[nnmb].index
 189                .. decoration
 190                .. '\n'
 191          )
 192       end
 193
 194       file:write('\n')
 195    end
 196
 197    file:write('}\n')
 198
 199 end
 200
 201 ----------------------------------------------------------------------
 202
 203 function DAG:updateOutput(input)
 204    self:putInOrder()
 205
 206    self:nestedApply(
 207       function(nnm, i)
 208          self.node[nnm].input = i
 209          -- nnm:updateOutput(i)
 210          self:rethrowErrors(nnm, self.node[nnm].index, 'updateOutput', i)
 211       end,
 212       self.inputModules,
 213       input
 214    )
 215
 216    for _, nnm in ipairs(self.sorted) do
 217       local node = self.node[nnm]
 218       if #node.pred > 0 then
 219          local i
 220          if #node.pred == 1 then
 221             i = node.pred[1].output
 222          elseif #node.pred > 1 then
 223             i = {}
 224             for k = 1, #node.pred do
 225                i[k] = node.pred[k].output
 226             end
 227          end
 228          node.input = i
 229          -- nnm:updateOutput(i)
 230          self:rethrowErrors(nnm, self.node[nnm].index, 'updateOutput', i)
 231       end
 232    end
 233
 234    self.output = self:nestedApply(
 235       function(m) return m.output end,
 236       self.outputModules
 237    )
 238
 239    return self.output
 240 end
 241
 242 function DAG:updateGradInput(input, gradOutput)
 243    assert(self.sorted, 'There has been a DAG structure change before a DAG:updateGradInput')
 244
 245    self:nestedApply(
 246       function(nnm, go)
 247          local node = self.node[nnm]
 248          node.gradOutput = go
 249          -- nnm:updateGradInput(self.node[nnm].input, go)
 250          self:rethrowErrors(nnm, node.index, 'updateGradInput', self.node[nnm].input, go)
 251       end,
 252       self.outputModules, gradOutput
 253    )
 254
 255    self:nestedApply(
 256       function(nnm, i) self.node[nnm].input = i end,
 257       self.inputModules, input
 258    )
 259
 260    for _, node in pairs(self.node) do
 261       node.gradInputSucc = {}
 262    end
 263
 264    for k = #self.sorted, 1, -1 do
 265       local nnm = self.sorted[k]
 266       local node = self.node[nnm]
 267       local pred, gradInputSucc = node.pred, node.gradInputSucc
 268
 269       if #gradInputSucc > 0 then
 270          node.gradOutput = self:computeGradOutput(gradInputSucc)
 271          -- nnm:updateGradInput(node.input, node.gradOutput)
 272          self:rethrowErrors(nnm, self.node[nnm].index, 'updateGradInput', node.input, node.gradOutput)
 273       end
 274
 275       -- We fill the gradInputSucc of our predecessors
 276       if #pred == 1 then
 277          table.insert(self.node[pred[1]].gradInputSucc, nnm.gradInput)
 278       elseif #pred > 1 then
 279          if not torch.type(nnm.gradInput) == 'table' then
 280             error('Should have a table gradInput since it has multiple predecessors')
 281          end
 282          for n = 1, #pred do
 283             table.insert(self.node[node.pred[n]].gradInputSucc, nnm.gradInput[n])
 284          end
 285       end
 286    end
 287
 288    self.gradInput = self:nestedApply(function(m) return m.gradInput end, self.inputModules)
 289
 290    return self.gradInput
 291 end
 292
 293 function DAG:accGradParameters(input, gradOutput, scale)
 294    scale = scale or 1
 295
 296    assert(self.sorted, 'There has been a DAG structure change before a DAG:accGradParameters')
 297
 298    self:nestedApply(
 299       function(nnm, go) self.node[nnm].gradOutput = go end,
 300       self.outputModules, gradOutput
 301    )
 302
 303    self:nestedApply(
 304       function(nnm, i) self.node[nnm].input = i end,
 305       self.inputModules, input
 306    )
 307
 308    for k = 1, #self.modules do
 309       local nnm = self.modules[k]
 310       local node = self.node[nnm]
 311       -- nnm:accGradParameters(node.input, node.gradOutput, scale)
 312       self:rethrowErrors(nnm, k, 'accGradParameters', node.input, node.gradOutput, scale)
 313    end
 314 end