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