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    for _, nnm in pairs(self.modules) do
  80       assert(distance[nnm], 'Some modules are not connected to inputs')
  81    end
  82
  83    self.sorted = {}
  84    for m, d in pairs(distance) do
  85       table.insert(self.sorted, { distance = d, nnm = m })
  86    end
  87
  88    table.sort(self.sorted, function(a, b) return a.distance < b.distance end)
  89
  90    for i, a in ipairs(self.sorted) do self.sorted[i] = a.nnm end
  91 end
  92
  93 -- This accumulates x in a where they are both nested tables of
  94 -- tensors. If first is true, set a = x. Behavior is undefined if a
  95 -- and x do not have the exact same structure.
  96 function DAG:nestedAccTensor(a, x, first)
  97    if torch.type(x) == 'table' then
  98       local b = {}
  99       for i in pairs(x) do
 100          b[i] = self:nestedAccTensor(a[i], x[i], first)
 101       end
 102       a = b
 103    else
 104       if first then
 105          if a then
 106             a:resizeAs(x):copy(x)
 107          else
 108             a = x:clone()
 109          end
 110       else
 111          a:add(x)
 112       end
 113    end
 114    return a
 115 end
 116
 117 function DAG:updateGradOutput(node)
 118    local gradInputSucc = node.gradInputSucc
 119    if #gradInputSucc == 1 then
 120       node.gradOutput = gradInputSucc[1]
 121    elseif #gradInputSucc > 1 then
 122       for k = 1, #gradInputSucc do
 123          node.gradOutput = self:nestedAccTensor(node.gradOutput, gradInputSucc[k], k == 1)
 124       end
 125    end
 126 end
 127
 128 ----------------------------------------------------------------------
 129
 130 -- Connect a sequence of modules
 131 function DAG:connect(...)
 132    self.sorted = nil
 133    local prev
 134    for _, nnm in pairs({...}) do
 135       self:createNode(nnm)
 136       if prev then
 137          table.insert(self.node[nnm].pred, prev)
 138          table.insert(self.node[prev].succ, nnm)
 139       end
 140       prev = nnm
 141    end
 142 end
 143
 144 function DAG:setInput(i)
 145    self.sorted = nil
 146    self.inputModules = i
 147    self:nestedApply(
 148       function(nnm)
 149          if #self.node[nnm].succ == 0 then
 150             error('Input modules must have outgoing edges.')
 151          end
 152          if #self.node[nnm].pred > 0 then
 153             error('Input modules cannot have incoming edges.')
 154          end
 155       end,
 156       self.inputModules
 157    )
 158 end
 159
 160 function DAG:setOutput(o)
 161    self.sorted = nil
 162    self.outputModules = o
 163    self:nestedApply(
 164       function(nnm)
 165          if #self.node[nnm].pred == 0 then
 166             error('Output module must have incoming edges.')
 167          end
 168          if #self.node[nnm].succ > 0 then
 169             error('Output module cannot have outgoing edges.')
 170          end
 171       end,
 172       self.outputModules
 173    )
 174 end
 175
 176 function DAG:print()
 177    self:putInOrder()
 178
 179    for i, d in ipairs(self.sorted) do
 180       print('#' .. i .. ' -> ' .. torch.type(d))
 181    end
 182 end
 183
 184 ----------------------------------------------------------------------
 185
 186 function DAG:saveDot(filename)
 187    local file = (filename and io.open(filename, 'w')) or io.stdout
 188
 189    file:write('digraph {\n')
 190
 191    file:write('\n')
 192
 193    for nnmb, node in pairs(self.node) do
 194       file:write(
 195          '  '
 196             .. node.index
 197             .. ' [shape=box,label=\"' .. torch.type(nnmb) .. '\"]'
 198             .. '\n'
 199       )
 200
 201       for i, nnma in pairs(node.pred) do
 202          local decoration = ''
 203          if #node.pred > 1 then
 204             -- decoration = ' [headlabel=\"' .. i .. '\"]'
 205             decoration = ' [label=\"' .. i .. '\"]'
 206          end
 207          file:write(
 208             '  '
 209                .. self.node[nnma].index
 210                .. ' -> '
 211                .. self.node[nnmb].index
 212                .. decoration
 213                .. '\n'
 214          )
 215       end
 216
 217       file:write('\n')
 218    end
 219
 220    file:write('}\n')
 221
 222 end
 223
 224 ----------------------------------------------------------------------
 225
 226 function DAG:updateOutput(input)
 227    self:putInOrder()
 228
 229    self:nestedApply(
 230       function(nnm, i)
 231          local node = self.node[nnm]
 232          node.input = i
 233          self:rethrowErrors(nnm, node.index, 'updateOutput', i)
 234       end,
 235       self.inputModules,
 236       input
 237    )
 238
 239    for _, nnm in ipairs(self.sorted) do
 240       local node = self.node[nnm]
 241       if #node.pred > 0 then
 242          local i
 243          if #node.pred == 1 then
 244             i = node.pred[1].output
 245          elseif #node.pred > 1 then
 246             i = {}
 247             for k = 1, #node.pred do
 248                i[k] = node.pred[k].output
 249             end
 250          end
 251          node.input = i
 252          self:rethrowErrors(nnm, node.index, 'updateOutput', i)
 253       end
 254    end
 255
 256    self.output = self:nestedApply(
 257       function(m) return m.output end,
 258       self.outputModules
 259    )
 260
 261    return self.output
 262 end
 263
 264 function DAG:updateGradInput(input, gradOutput)
 265    assert(self.sorted, 'There has been a DAG structure change before a DAG:updateGradInput')
 266
 267    self:nestedApply(
 268       function(nnm, go)
 269          local node = self.node[nnm]
 270          node.gradOutput = go
 271          self:rethrowErrors(nnm, node.index, 'updateGradInput', node.input, go)
 272       end,
 273       self.outputModules, gradOutput
 274    )
 275
 276    self:nestedApply(
 277       function(nnm, i) self.node[nnm].input = i end,
 278       self.inputModules, input
 279    )
 280
 281    for _, node in pairs(self.node) do
 282       node.gradInputSucc = {}
 283    end
 284
 285    for k = #self.sorted, 1, -1 do
 286       local nnm = self.sorted[k]
 287       local node = self.node[nnm]
 288       local pred = node.pred
 289
 290       if #node.gradInputSucc > 0 then
 291          self:updateGradOutput(node)
 292          self:rethrowErrors(nnm, node.index, 'updateGradInput', node.input, node.gradOutput)
 293       end
 294
 295       -- We fill the gradInputSucc of our predecessors
 296       if #pred == 1 then
 297          table.insert(self.node[pred[1]].gradInputSucc, nnm.gradInput)
 298       elseif #pred > 1 then
 299          if not torch.type(nnm.gradInput) == 'table' then
 300             error('Should have a table gradInput since it has multiple predecessors')
 301          end
 302          for n = 1, #pred do
 303             table.insert(self.node[node.pred[n]].gradInputSucc, nnm.gradInput[n])
 304          end
 305       end
 306    end
 307
 308    self.gradInput = self:nestedApply(function(m) return m.gradInput end, self.inputModules)
 309
 310    return self.gradInput
 311 end
 312
 313 function DAG:accGradParameters(input, gradOutput, scale)
 314    assert(self.sorted, 'There has been a DAG structure change before a DAG:accGradParameters')
 315
 316    self:nestedApply(
 317       function(nnm, go) self.node[nnm].gradOutput = go end,
 318       self.outputModules, gradOutput
 319    )
 320
 321    self:nestedApply(
 322       function(nnm, i) self.node[nnm].input = i end,
 323       self.inputModules, input
 324    )
 325
 326    for k = 1, #self.modules do
 327       local nnm = self.modules[k]
 328       local node = self.node[nnm]
 329       self:rethrowErrors(nnm, k, 'accGradParameters', node.input, node.gradOutput, scale)
 330    end
 331 end
 332
 333 function DAG:clearState()
 334    self.sorted = nil
 335    for _, node in pairs(self.node) do
 336       node.gradInputSucc = nil
 337       node.input = nil
 338       node.gradOutput = nil
 339    end
 340    return parent.clearState(self)
 341 end