fix implicit methods

Author: sungyubkim

gbml/imaml.py

@@ -14,16 +14,15 @@ def __init__(self, args):
         super().__init__(args)
         self._init_net()
         self._init_opt()
-        self.lamb = 2.0
-        self.n_cg = 3
+        self.lamb = 100.0
+        self.n_cg = 1
         return None
 
     @torch.enable_grad()
     def inner_loop(self, fmodel, diffopt, train_input, train_target):
 
         train_logit = fmodel(train_input)
         inner_loss = F.cross_entropy(train_logit, train_target)
-        inner_loss += (self.lamb/2.) * ((torch.nn.utils.parameters_to_vector(self.network.parameters())-torch.nn.utils.parameters_to_vector(self.network.parameters()).detach())**2).sum()
         diffopt.step(inner_loss)
 
         return None
@@ -41,8 +40,6 @@ def cg(self, in_grad, outer_grad, params):
             beta = (r_new @ r_new)/(r @ r)
             p = r_new + beta * p
             r = r_new.clone().detach()
-        #     print(alpha, beta ,r @ r);input()
-        # print('end')
         return self.vec_to_grad(x)
 
     def vec_to_grad(self, vec):

gbml/neumann.py

@@ -14,7 +14,8 @@ def __init__(self, args):
         super().__init__(args)
         self._init_net()
         self._init_opt()
-        self.n_series = 3
+        self.lamb = 100.0
+        self.n_series = 1
         return None
 
     @torch.enable_grad()
@@ -48,7 +49,7 @@ def vec_to_grad(self, vec):
     def hv_prod(self, in_grad, x, params):
         hv = torch.autograd.grad(in_grad, params, retain_graph=True, grad_outputs=x)
         hv = torch.nn.utils.parameters_to_vector(hv)
-        hv = (-1.*self.args.inner_lr) * hv # scale for regularization
+        hv = (-1./self.lamb) * hv # scaling for convergence
         return hv.detach()
 
     def outer_loop(self, batch, is_train):