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