Added diag hessian support for a few modules.

1 files changed, 13 insertions, 0 deletions

diff --git a/init.lua b/init.lua
index aff1468..5f3daf2 100644
--- a/init.lua
+++ b/init.lua
@@ -108,6 +108,7 @@ torch.include('nnx', 'SGDOptimization.lua')
 torch.include('nnx', 'LBFGSOptimization.lua')
 torch.include('nnx', 'CGOptimization.lua')
 torch.include('nnx', 'GeneticSGDOptimization.lua')
+torch.include('nnx', 'DiagHessian.lua')
 
 -- trainers:
 torch.include('nnx', 'Trainer.lua')
@@ -196,6 +197,18 @@ function nnx.getGradParameters(...)
    return holder
 end
 
+function nnx.getDiagHessianParameters(...)
+   -- to hold all parameters found
+   local holder = {}
+   -- call recursive call
+   local modules = {...}
+   for _,module in ipairs(modules) do
+      get(module, holder, {'diagHessianWeight', 'diagHessianBias'})
+   end
+   -- return all parameters found
+   return holder
+end
+
 function nnx.flattenParameters(parameters)
    -- already flat ?
    local flat = true