def matrices_from_kronecker_curvature(self, extension_cls, savefield): results = [] with backpack(extension_cls()): self.loss().backward() for p in self.model.parameters(): factors = getattr(p, savefield) results.append(kfacs_to_mat(factors)) return results
def test_kfacs_to_mat(self): """Check matrix from list of Kronecker factors with `scipy`.""" for _ in range(self.RUNS): factors = self.make_random_kfacs() bp_result = bp_utils.kfacs_to_mat(factors) sp_result = self.scipy_kfacs_to_mat(factors) assert self.allclose(bp_result, sp_result)
def test_kfac_should_approx_ggn_montecarlo(problem: ExtensionsTestProblem): """Check that for batch_size = 1, the K-FAC is the same as the GGN. Should be true for linear layers and in the limit of infinite mc_samples. Args: problem: Test case. """ problem.set_up() autograd_res = AutogradExtensions(problem).ggn_blocks() mc_samples = 300000 backpack_kfac = BackpackExtensions(problem).kfac_chunk(mc_samples) backpack_res = [kfacs_to_mat(kfac) for kfac in backpack_kfac] check_sizes_and_values(autograd_res, backpack_res, atol=5e-3, rtol=5e-3) problem.tear_down()
def hbp_with_curv( self, curv_type, loss_hessian_strategy=LossHessianStrategy.SUM, backprop_strategy=BackpropStrategy.BATCH_AVERAGE, ea_strategy=ExpectationApproximation.BOTEV_MARTENS, ): results = [] with backpack( new_ext.HBP( curv_type=curv_type, loss_hessian_strategy=loss_hessian_strategy, backprop_strategy=backprop_strategy, ea_strategy=ea_strategy, )): self.loss().backward() for p in self.model.parameters(): factors = p.hbp results.append(kfacs_to_mat(factors)) return results
def set_up(): factors = self.make_random_kfacs() kfac = bp_utils.kfacs_to_mat(factors) tensor = make_tensor(kfac) return factors, kfac, tensor