dagnn.lua

   1
   2 require 'torch'
   3 require 'nn'
   4
   5 local DAG, parent = torch.class('nn.DAG', 'nn.Container')
   6
   7 function DAG:__init()
   8    parent.__init(self)
   9    self.pred = {}
  10    self.succ = {}
  11 end
  12
  13 function DAG:addEdge(a, b)
  14    self.sorted = nil
  15    local pred, succ = self.pred, self.succ
  16    if not pred[a] and not succ[a] then
  17       self:add(a)
  18    end
  19    if not pred[b] and not succ[b] then
  20       self:add(b)
  21    end
  22    pred[b] = pred[b] or {}
  23    pred[b][#pred[b] + 1] = a
  24    succ[a] = succ[a] or {}
  25    succ[a][#succ[a] + 1] = b
  26 end
  27
  28 function DAG:setInput(i)
  29    self.sorted = nil
  30    if torch.type(i) == 'table' then
  31       self.inputModules = i
  32       for _, m in ipairs(i) do
  33          if not self.pred[m] and not self.succ[m] then
  34             self:add(m)
  35          end
  36       end
  37    else
  38       self:setInput({ i })
  39    end
  40 end
  41
  42 function DAG:setOutput(o)
  43    self.sorted = nil
  44    if torch.type(o) == 'table' then
  45       self.outputModules = o
  46       for _, m in ipairs(o) do
  47          if not self.pred[m] and not self.succ[m] then
  48             self:add(m)
  49          end
  50       end
  51    else
  52       self:setOutput({ o })
  53    end
  54 end
  55
  56 function DAG:sort()
  57    if self.sorted then
  58       return
  59    end
  60
  61    local distance = {}
  62
  63    for _, a in pairs(self.inputModules) do
  64       distance[a] = 1
  65    end
  66
  67    local nc
  68
  69    repeat
  70       nc = 0
  71       for i, isucc in pairs(self.succ) do
  72          for _, j in pairs(isucc) do
  73             if distance[i] and (not distance[j] or distance[j] < distance[i] + 1) then
  74                distance[j] = distance[i] + 1
  75                nc = nc + 1
  76             end
  77          end
  78       end
  79    until nc == 0
  80
  81    self.sorted = { }
  82    for i, d in pairs(distance) do
  83       table.insert(self.sorted, { d, i })
  84    end
  85
  86    table.sort(self.sorted, function(a, b) return a[1] < b[1] end)
  87    for i, a in ipairs(self.sorted) do self.sorted[i] = a[2] end
  88 end
  89
  90 function DAG:print()
  91    self:sort()
  92
  93    for i, d in ipairs(self.sorted) do
  94       print('#' .. i .. ' -> ' .. torch.type(d))
  95    end
  96 end
  97
  98 function DAG:updateOutput(input)
  99    self:sort()
 100
 101    if #self.inputModules == 1 then
 102       self.inputModules[1]:updateOutput(input)
 103    else
 104       for i, d in ipairs(self.inputModules) do
 105          d:updateOutput(input[i])
 106       end
 107    end
 108
 109    for _, d in ipairs(self.sorted) do
 110       if self.pred[d] then
 111          if #self.pred[d] == 1 then
 112             d:updateOutput(self.pred[d][1].output)
 113          elseif #self.pred[d] > 1 then
 114             local c = {}
 115             for k = 1, #self.pred[d] do
 116                c[k] = self.pred[d][k].output
 117             end
 118             d:updateOutput(c)
 119          end
 120       end
 121    end
 122
 123    if #self.outputModules == 1 then
 124       self.output = self.outputModules[1].output
 125    else
 126       self.output = { }
 127       for i, d in ipairs(self.outputModules) do
 128          self.output[i] = d.output
 129       end
 130    end
 131
 132    return self.output
 133 end
 134
 135 function DAG:updateGradInput(input, gradOutput)
 136    self:sort()
 137 end
 138
 139 return DAG