ncdu-model-explore

json = require 'cjson'

function buildNcduLayer(name, module)
   local result = nil
   if torch.isTensor(module) then
      if module:numel() ~= 0 then
         local strt = {name..': [' .. torch.typename(module) .. ' of size '}
         for i=1,module:nDimension() do
            table.insert(strt, module:size(i))
            if i ~= module:nDimension() then
               table.insert(strt, 'x')
            end
         end
         table.insert(strt, ']')
         result = {name = table.concat(strt),
                   dsize = module:numel() * module:storage():elementSize()
                }
      else
         result = {name = name..": [empty "..torch.typename(module).."]"}
      end
   elseif type(module)=="table" and module.modules then
      result = { {name = name..": "..string.gsub(tostring(module), "\n", " ")} }
      for i,m in ipairs(module.modules) do
      
         table.insert(result, buildNcduLayer(string.format("%03d", i), m))
      end
   elseif type(module)=="table" then
      result = {{name=name..": "..string.gsub(tostring(module), "\n", " ")}}
      for k,v in pairs(module) do
         table.insert(result, buildNcduLayer(k,v))
      end
   else
      result = {name=name.." (primitive)", dsize=0}
   end
   return result
end

function buildNcdu(model)
   local result = {1,0,{timestamp=1451677436,progver="0.1",progname="ncdu-model-explore"}}
   table.insert(result, buildNcduLayer("model", model))
   return json.encode(result)
end

function exploreNcdu(model)
   local tmpname = os.tmpname()
   local tmphandle = io.open(tmpname, "w")
   
   tmphandle:write(buildNcdu(model))
   tmphandle:close()
   os.execute("ncdu -f "..tmpname)
   os.remove(tmpname)
end