def hmp(self, mat_list):
     assert len(mat_list) == len(list(self.model.parameters()))
     results = []
     with backpack(new_ext.HMP()):
         self.loss().backward()
         for p, mat in zip(self.model.parameters(), mat_list):
             results.append(p.hmp(mat))
     return results
Example #2
0
def test_interface_hmp_conv():
    interface_test(new_ext.HMP(), use_conv=True)
Example #3
0
def test_interface_hmp():
    interface_test(new_ext.HMP())
Example #4
0
        return CGNOptimizer(
            model.parameters(),
            extension,
            lr=LR,
            damping=DAMPING,
            maxiter=CG_MAX_ITER,
            tol=CG_TOL,
            atol=CG_ATOL,
        )

    return optimizer_fn


curvatures = [
    extensions.GGNMP(),
    extensions.HMP(),
    extensions.PCHMP(modify="abs"),
    extensions.PCHMP(modify="clip"),
]

labels = [
    "GGN",
    "Hessian",
    "PCH-abs",
    "PCH-clip",
]

optimizers = []
for curvature in curvatures:
    optimizers.append(make_cgn_optimizer_fn(curvature))