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