From 941ae69a9ccf274e3bff8b1fa45da04fb825a647 Mon Sep 17 00:00:00 2001 From: Francois Fleuret Date: Thu, 12 Jan 2017 14:48:02 +0100 Subject: [PATCH] Update. --- dagnn.lua | 148 ++++++++++++++++++++++++++++--------------------- test-dagnn.lua | 11 +++- 2 files changed, 93 insertions(+), 66 deletions(-) diff --git a/dagnn.lua b/dagnn.lua index 1b467e7..0b8f7d4 100755 --- a/dagnn.lua +++ b/dagnn.lua @@ -6,34 +6,36 @@ local DAG, parent = torch.class('nn.DAG', 'nn.Container') function DAG:__init() parent.__init(self) - self.pred = {} - self.succ = {} + -- Nodes are indexed by the module they encompass + self.node = { } +end + +function DAG:createNode(n) + if not self.node[n] then + self:add(n) -- Add it to the object as a Container + self.node[n] = {} + self.node[n].succ = {} + self.node[n].pred = {} + end end function DAG:addEdge(a, b) self.sorted = nil - local pred, succ = self.pred, self.succ - if not pred[a] and not succ[a] then - self:add(a) - end - if not pred[b] and not succ[b] then - self:add(b) - end - pred[b] = pred[b] or {} - pred[b][#pred[b] + 1] = a - succ[a] = succ[a] or {} - succ[a][#succ[a] + 1] = b + self:createNode(a) + self:createNode(b) + table.insert(self.node[b].pred, a) + table.insert(self.node[a].succ, b) end -- Apply f on t recursively; use the corresponding a1 and a2 elements -- (i.e. same keys) as second and third parameters to f when -- available; return the results from f, organized in a similarly -- nested table. -function DAG:applyOnModules(f, t, a1, a2) +function DAG:nestApply(f, t, a1, a2) if torch.type(t) == 'table' then local result = {} for k, s in pairs(t) do - result[k] = self:applyOnModules(f, s, a1 and a1[k], a2 and a2[k]) + result[k] = self:nestApply(f, s, a1 and a1[k], a2 and a2[k]) end return result else @@ -44,12 +46,12 @@ end function DAG:setInput(i) self.sorted = nil self.inputModules = i - self:applyOnModules( + self:nestApply( function(m) - if not self.succ[m] or #self.succ[m] == 0 then + if #self.node[m].succ == 0 then error('Input modules must have outgoing edges.') end - if self.pred[m] and #self.pred[m] > 0 then + if #self.node[m].pred > 0 then error('Input modules cannog have incoming edges.') end end, @@ -60,12 +62,12 @@ end function DAG:setOutput(o) self.sorted = nil self.outputModules = o - self:applyOnModules( + self:nestApply( function(m) - if not self.pred[m] or #self.pred[m] == 0 then + if #self.node[m].pred == 0 then error('Output module must have incoming edges.') end - if self.succ[m] and #self.succ[m] > 0 then + if #self.node[m].succ > 0 then error('Output module cannot have outgoing edges.') end end, @@ -73,21 +75,23 @@ function DAG:setOutput(o) ) end -function DAG:sort() +function DAG:putInOrder() if self.sorted then return end + -- First, we sort the nodes according to the DAG order + local distance = {} - self:applyOnModules(function(m) distance[m] = 1 end, self.inputModules) + self:nestApply(function(m) distance[m] = 1 end, self.inputModules) local nc repeat nc = 0 - for i, isucc in pairs(self.succ) do - for _, j in pairs(isucc) do + for i, node in pairs(self.node) do + for _, j in pairs(node.succ) do if distance[i] and (not distance[j] or distance[j] < distance[i] + 1) then distance[j] = distance[i] + 1 nc = nc + 1 @@ -97,16 +101,17 @@ function DAG:sort() until nc == 0 self.sorted = { } - for i, d in pairs(distance) do - table.insert(self.sorted, { d, i }) + for n, d in pairs(distance) do + table.insert(self.sorted, { distance = d, node = n }) end - table.sort(self.sorted, function(a, b) return a[1] < b[1] end) - for i, a in ipairs(self.sorted) do self.sorted[i] = a[2] end + table.sort(self.sorted, function(a, b) return a.distance < b.distance end) + + for i, a in ipairs(self.sorted) do self.sorted[i] = a.node end end function DAG:print() - self:sort() + self:putInOrder() for i, d in ipairs(self.sorted) do print('#' .. i .. ' -> ' .. torch.type(d)) @@ -114,57 +119,74 @@ function DAG:print() end function DAG:updateOutput(input) - self:sort() - - self:applyOnModules(function(m, i) m:updateOutput(i) end, self.inputModules, input) - - for _, d in ipairs(self.sorted) do - if self.pred[d] then - if #self.pred[d] == 1 then - d:updateOutput(self.pred[d][1].output) - elseif #self.pred[d] > 1 then - local c = {} - for k = 1, #self.pred[d] do - c[k] = self.pred[d][k].output + self:putInOrder() + + self:nestApply(function(m, i) m:updateOutput(i) end, self.inputModules, input) + + for _, m in ipairs(self.sorted) do + if #self.node[m].pred > 0 then + local i + if #self.node[m].pred == 1 then + i = self.node[m].pred[1].output + elseif #self.node[m].pred > 1 then + i = {} + for k = 1, #self.node[m].pred do + i[k] = self.node[m].pred[k].output end - d:updateOutput(c) end + self.node[m].input = i + m:updateOutput(i) end end - self.output = self:applyOnModules(function(m) return m.output end, self.outputModules) + self.output = self:nestApply(function(m) return m.output end, self.outputModules) return self.output end function DAG:updateGradInput(input, gradOutput) - self:sort() + self:putInOrder() - self:applyOnModules( - function(m, i, go) m:updateGradInput(i, go) end, - self.outputModules, input, gradOutput + self:nestApply( + function(m, go) m:updateGradInput(self.node[m].input, go) end, + self.outputModules, gradOutput ) - for k = self.sorted, 1, -1 do - local m = sorted[k] - if self.succ[d] then - if #self.succ[d] == 1 then - d:updateGradInput(self.succ[d][1].gradInput) - elseif #self.succ[d] > 1 then - local sum - for k = 1, #self.succ[d] do - if sum then - sum:add(self.succ[d][k].gradInput) - else - sum = self.succ[d][k].gradInput:clone() - end + for _, node in pairs(self.node) do + node.gradInputSucc = {} + end + + for k = #self.sorted, 1, -1 do + local m = self.sorted[k] + local node = self.node[m] + local pred, succ, gradInputSucc = node.pred, node.succ, node.gradInputSucc + + -- We update m:gradInput + if #gradInputSucc == 1 then + m:updateGradInput(node.input, gradInputSucc[1]) + elseif #gradInputSucc > 1 then + local sum + for k = 1, #succ do + if sum then + sum:add(succ[k].gradInput) + else + sum = succ[k].gradInput end - d:updateGradInput(sum) + end + m:updateGradInput(node.input, sum) + end + + -- We fill the gradInputSucc of our predecessors + if #pred == 1 then + table.insert(self.node[pred[1]].gradInputSucc, node.gradInput) + elseif #pred > 1 then + for n = 1, #pred do + table.insert(self.node[node.pred[n]].gradInputSucc, m.gradInput[n]) end end end - self.gradInput = self:applyOnModules(function(m) return m.gradInput end, self.inputModules) + self.gradInput = self:nestApply(function(m) return m.gradInput end, self.inputModules) return self.gradInput end diff --git a/test-dagnn.lua b/test-dagnn.lua index 434f662..0c9fe6d 100755 --- a/test-dagnn.lua +++ b/test-dagnn.lua @@ -42,17 +42,22 @@ g:addEdge(b, c) g:addEdge(b, d) g:addEdge(d, f) -g:setInput({a}) -g:setOutput({e, f}) +g:setInput({ a }) +g:setOutput({ e, f }) g:print() input = torch.Tensor(3, 10):uniform() -output = g:updateOutput({input}) +output = g:updateOutput({ input }) printTensorTable(output) ---------------------------------------------------------------------- +print('******************************************************************') +print('** updateGradInput ***********************************************') +print('******************************************************************') gradInput = g:updateGradInput({ input }, output) + +printTensorTable(gradInput) -- 2.20.1