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
def test_interface_hmp_conv():
interface_test(new_ext.HMP(), use_conv=True)
def test_interface_hmp():
interface_test(new_ext.HMP())
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))