diff options
author | Ivo Danihelka <danihelka@google.com> | 2015-04-01 16:22:21 +0300 |
---|---|---|
committer | Ivo Danihelka <danihelka@google.com> | 2015-04-01 16:22:21 +0300 |
commit | 2d308e6a1457f89cdfb3207921d0322f5389b2a9 (patch) | |
tree | c887e3c6ad005f8fd9871561c3f3e32e69007866 | |
parent | 950c67076477b70ecf904ce899fda98fbada965d (diff) |
Allowed to annotate nodes based on the local variable names.
-rw-r--r-- | graphinspecting.lua | 20 |
1 files changed, 20 insertions, 0 deletions
diff --git a/graphinspecting.lua b/graphinspecting.lua index 045be31..66a334f 100644 --- a/graphinspecting.lua +++ b/graphinspecting.lua @@ -117,3 +117,23 @@ function nngraph.setDebug(enable) end end end + +-- Sets node.data.annotations.name for the found nodes. +-- The local variables at the given stack level are inspected. +-- The default stack level is 1 (the function that called annotateNodes()). +function nngraph.annotateNodes(stackLevel) + stackLevel = stackLevel or 1 + for index = 1, math.huge do + local varName, varValue = debug.getlocal(stackLevel + 1, index) + if not varName then + break + end + if torch.typename(varValue) == "nngraph.Node" then + -- An explicit name is preserved. + if not varValue.data.annotations.name then + varValue:annotate({name = varName}) + end + end + end +end + |