dagnn.lua

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