From: Francois Fleuret Date: Mon, 16 Jan 2017 06:39:34 +0000 (+0100) Subject: Added DAG:setLabel to add a label to a module in the graph. X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=commitdiff_plain;h=cd1aa6aaa75f5d5b281be0cfbacd51991b3b1ca3;p=dagnn.git Added DAG:setLabel to add a label to a module in the graph. --- diff --git a/dagnn.lua b/dagnn.lua index ca26926..f9d6ff9 100755 --- a/dagnn.lua +++ b/dagnn.lua @@ -148,6 +148,10 @@ function DAG:connect(...) end end +function DAG:setLabel(nnm, label) + self.node[nnm].label = label +end + function DAG:setInput(i) self.sorted = nil self.inputModules = i @@ -176,7 +180,11 @@ function DAG:print() self:putInOrder() for i, d in ipairs(self.sorted) do - print('#' .. i .. ' -> ' .. torch.type(d)) + local decoration = '' + if self.node[d].label then + decoration = ' [' .. self.node[d].label .. ']' + end + print('#' .. i .. ' -> ' .. torch.type(d) .. decoration) end end @@ -211,7 +219,7 @@ function DAG:saveDot(filename) file:write( ' ' .. node.index - .. ' [shape=box,label=\"' .. torch.type(nnmb) .. '\"]' + .. ' [shape=box,label=\"' .. (self.node[nnmb].label or torch.type(nnmb)) .. '\"]' .. '\n' )