X-Git-Url: https://www.fleuret.org/cgi-bin/gitweb/gitweb.cgi?p=dagnn.git;a=blobdiff_plain;f=dagnn.lua;fp=dagnn.lua;h=ca2692692379f7d19db4aa82e7c6af7a48bbd629;hp=1f45b2acae3bc546495e2b33bd0666f6f577c54e;hb=78671207db483567021a935d0738ba85d8b16551;hpb=f3a748c4119c4f374b30d0010198732bb80b43a6

diff --git a/dagnn.lua b/dagnn.lua
index 1f45b2a..ca26926 100755
--- a/dagnn.lua
+++ b/dagnn.lua
@@ -185,10 +185,28 @@ end
 function DAG:saveDot(filename)
    local file = (filename and io.open(filename, 'w')) or io.stdout
 
+   local function writeNestedCluster(prefix, list, indent)
+      local indent = indent or ''
+      if torch.type(list) == 'table' then
+         file:write(indent .. '  subgraph cluster_' .. prefix .. ' {\n');
+         for k, x in pairs(list) do
+            writeNestedCluster(prefix .. '_' .. k, x, '  ' .. indent)
+         end
+         file:write(indent .. '  }\n');
+      else
+         file:write(indent .. '  ' .. self.node[list].index .. ' [color=red]\n')
+      end
+   end
+
    file:write('digraph {\n')
 
    file:write('\n')
 
+   writeNestedCluster('input', self.inputModules)
+   writeNestedCluster('output', self.outputModules)
+
+   file:write('\n')
+
    for nnmb, node in pairs(self.node) do
       file:write(
          '  '