X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=dagnn.lua;h=ca51841f8442417d10097e6a6939fd86f3bceda2;hb=d8fd868f94ce0b66cd2cc1a4615df10a88b5d5ec;hp=158ef785d446fb22f6226248c6a07c83cf47b6c3;hpb=e6516772e13dd5424f0a1b7e2063a7417614844c;p=dagnn.git diff --git a/dagnn.lua b/dagnn.lua index 158ef78..ca51841 100755 --- a/dagnn.lua +++ b/dagnn.lua @@ -29,24 +29,7 @@ function DAG:__init() self.node = { } end -function DAG:createNode(nnm) - if not self.node[nnm] then - self:add(nnm) -- Add it to the object as a Container - self.node[nnm] = {} - self.node[nnm].succ = {} - self.node[nnm].pred = {} - end -end - -function DAG:addEdge(nnma, nnmb) - self.sorted = nil - self:createNode(nnma) - self:createNode(nnmb) - table.insert(self.node[nnmb].pred, nnma) - table.insert(self.node[nnma].succ, nnmb) -end - --- Apply f on t recursively; use the corresponding element from args +-- Apply f on t recursively; use the corresponding elements from args -- (i.e. same keys) as second parameter to f when available; return -- the results from f, organized in a similarly nested table. function DAG:nestedApply(f, t, args) @@ -61,6 +44,80 @@ function DAG:nestedApply(f, t, args) end end +function DAG:createNode(nnm) + if not self.node[nnm] then + self:add(nnm) -- Add it to the object as a Container + local node = {} + node.succ = {} + node.pred = {} + node.index = #self.modules + self.node[nnm] = node + end +end + +function DAG:putInOrder() + if self.sorted then + return + end + + local distance = {} + self:nestedApply(function(m) distance[m] = 1 end, self.inputModules) + + local nc + repeat + nc = 0 + for nnma, node in pairs(self.node) do + for _, nnmb in pairs(node.succ) do + if distance[nnma] and (not distance[nnmb] or distance[nnmb] < distance[nnma] + 1) then + distance[nnmb] = distance[nnma] + 1 + nc = nc + 1 + end + end + end + until nc == 0 + + self.sorted = { } + for m, d in pairs(distance) do + table.insert(self.sorted, { distance = d, nnm = m }) + 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.nnm end +end + +function DAG:computeGradOutput(gradInputSucc) + local gi + if #gradInputSucc == 1 then + gi = gradInputSucc[1] -- we avoid a clone() + elseif #gradInputSucc > 1 then + for k = 1, #gradInputSucc do + if gi then + gi:add(gradInputSucc[k]) + else + gi = gradInputSucc[k]:clone() + end + end + end + return gi +end + +---------------------------------------------------------------------- + +-- Connect a sequence of modules +function DAG:connect(...) + self.sorted = nil + local prev + for _, nnm in pairs({...}) do + self:createNode(nnm) + if prev then + table.insert(self.node[nnm].pred, prev) + table.insert(self.node[prev].succ, nnm) + end + prev = nnm + end +end + function DAG:setInput(i) self.sorted = nil self.inputModules = i @@ -93,48 +150,49 @@ function DAG:setOutput(o) ) end -function DAG:putInOrder() - if self.sorted then - return +function DAG:print() + self:putInOrder() + + for i, d in ipairs(self.sorted) do + print('#' .. i .. ' -> ' .. torch.type(d)) end +end - -- First, we sort the nodes according to the DAG order +---------------------------------------------------------------------- - local distance = {} +function DAG:saveDot(filename) + local file = (filename and io.open(filename, 'w')) or io.stdout - self:nestedApply(function(m) distance[m] = 1 end, self.inputModules) + file:write('digraph {\n') - local nc + file:write('\n') - repeat - nc = 0 - for nnma, node in pairs(self.node) do - for _, nnmb in pairs(node.succ) do - if distance[nnma] and (not distance[nnmb] or distance[nnmb] < distance[nnma] + 1) then - distance[nnmb] = distance[nnma] + 1 - nc = nc + 1 - end - end + for nnma, node in pairs(self.node) do + file:write( + ' ' + .. node.index + .. ' [shape=box,label=\"' .. torch.type(nnma) .. '\"]' + .. '\n' + ) + + for _, nnmb in pairs(node.succ) do + file:write( + ' ' + .. node.index + .. ' -> ' + .. self.node[nnmb].index + .. '\n' + ) end - until nc == 0 - self.sorted = { } - for m, d in pairs(distance) do - table.insert(self.sorted, { distance = d, nnm = m }) + file:write('\n') end - table.sort(self.sorted, function(a, b) return a.distance < b.distance end) + file:write('}\n') - for i, a in ipairs(self.sorted) do self.sorted[i] = a.nnm end end -function DAG:print() - self:putInOrder() - - for i, d in ipairs(self.sorted) do - print('#' .. i .. ' -> ' .. torch.type(d)) - end -end +---------------------------------------------------------------------- function DAG:updateOutput(input) self:putInOrder() @@ -142,7 +200,8 @@ function DAG:updateOutput(input) self:nestedApply( function(nnm, i) self.node[nnm].input = i - nnm:updateOutput(i) + -- nnm:updateOutput(i) + self:rethrowErrors(nnm, self.node[nnm].index, 'updateOutput', i) end, self.inputModules, input @@ -161,7 +220,8 @@ function DAG:updateOutput(input) end end node.input = i - nnm:updateOutput(i) + -- nnm:updateOutput(i) + self:rethrowErrors(nnm, self.node[nnm].index, 'updateOutput', i) end end @@ -173,27 +233,14 @@ function DAG:updateOutput(input) return self.output end -function DAG:computeGradInput(gradInputSucc) - local gi - if #gradInputSucc == 1 then - gi = gradInputSucc[1] -- we avoid a clone() - elseif #gradInputSucc > 1 then - for k = 1, #gradInputSucc do - if gi then - gi:add(gradInputSucc[k]) - else - gi = gradInputSucc[k]:clone() - end - end - end - return gi -end - function DAG:updateGradInput(input, gradOutput) - self:putInOrder() + assert(self.sorted, 'there has been a DAG structure change before a DAG:updateGradInput') self:nestedApply( - function(nnm, go) nnm:updateGradInput(self.node[nnm].input, go) end, + function(nnm, go) + -- nnm:updateGradInput(self.node[nnm].input, go) + self:rethrowErrors(nnm, self.node[nnm].index, 'updateGradInput', self.node[nnm].input, go) + end, self.outputModules, gradOutput ) @@ -212,7 +259,9 @@ function DAG:updateGradInput(input, gradOutput) local pred, gradInputSucc = node.pred, node.gradInputSucc if #gradInputSucc > 0 then - nnm:updateGradInput(node.input, self:computeGradInput(gradInputSucc)) + node.gradOutput = self:computeGradOutput(gradInputSucc) + -- nnm:updateGradInput(node.input, node.gradOutput) + self:rethrowErrors(nnm, self.node[nnm].index, 'updateGradInput', node.input, node.gradOutput) end -- We fill the gradInputSucc of our predecessors @@ -236,21 +285,12 @@ end function DAG:accGradParameters(input, gradOutput, scale) scale = scale or 1 - self:putInOrder() - - self:nestedApply( - function(nnm, go) nnm:updateGradInput(self.node[nnm].input, go) end, - self.outputModules, gradOutput - ) - - self:nestedApply( - function(nnm, i) self.node[nnm].input = i end, - self.inputModules, input - ) + assert(self.sorted, 'there has been a DAG structure change before a DAG:accGradParameters') - for k = #self.sorted, 1, -1 do - local nnm = self.sorted[k] + for k = 1, #self.modules do + local nnm = self.modules[k] local node = self.node[nnm] - nnm:accGradParameters(node.input, self:computeGradInput(node.gradInputSucc), scale) + -- nnm:accGradParameters(node.input, node.gradOutput, scale) + self:rethrowErrors(nnm, k, 'accGradParameters', node.input, self:computeGradOutput(node.gradInputSucc), scale) end end