graphnn.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.input = 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.output = 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.input) 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    return self.output.output
  96 end
  97
  98 ----------------------------------------------------------------------
  99
 100 a = nn.Linear(10, 10)
 101 b = nn.ReLU()
 102 c = nn.Linear(10, 3)
 103 d = nn.Linear(10, 3)
 104 e = nn.CMulTable()
 105
 106 --[[
 107
 108    a -----> b ---> c ---- e ---
 109              \           /
 110               \--> d ---/
 111
 112 ]]--
 113
 114 g = Graph:new()
 115
 116 g:setInput(a)
 117 g:setOutput(e)
 118 g:addEdge(c, e)
 119 g:addEdge(a, b)
 120 g:addEdge(d, e)
 121 g:addEdge(b, c)
 122 g:addEdge(b, d)
 123
 124 g:order()
 125 g:print(graph)