def test_ggn_mc(
problem: ExtensionsTestProblem, subsampling: Union[List[int], None]
) -> None:
"""Compare MC-approximated GGN from BackPACK with exact version from autograd.
Args:
problem: Test case with small network whose GGN can be evaluated.
subsampling: Indices of active samples. ``None`` uses the full mini-batch.
"""
skip_large_parameters(problem)
skip_subsampling_conflict(problem, subsampling)
autograd_res = AutogradExtensions(problem).ggn(subsampling=subsampling)
atol, rtol = 5e-3, 5e-3
mc_samples, chunks = 150000, 15
backpack_res = BackpackExtensions(problem).ggn_mc(
mc_samples, chunks=chunks, subsampling=subsampling
)
# compare normalized entries ∈ [-1; 1] (easier to tune atol)
max_val = max(autograd_res.abs().max(), backpack_res.abs().max())
# NOTE: The GGN can be exactly zero; e.g. if a ReLU after all parameters zeroes
# its input, its Jacobian is thus zero and will cancel the backpropagated GGN
if not isclose(max_val, 0):
autograd_res, backpack_res = autograd_res / max_val, backpack_res / max_val
check_sizes_and_values(autograd_res, backpack_res, atol=atol, rtol=rtol)
def test_batch_l2_grad_hook(problem):
"""Test squared ℓ₂ norm of individual gradients computed via extension hook.
Args:
problem (ExtensionsTestProblem): Problem for extension test.
"""
problem.set_up()
backpack_res = BackpackExtensions(problem).batch_l2_grad_extension_hook()
autograd_res = AutogradExtensions(problem).batch_l2_grad()
check_sizes_and_values(autograd_res, backpack_res)
problem.tear_down()
def test_sum_grad_squared_hook(problem):
"""Test individual gradient second moment computed via extension hook.
Args:
problem (ExtensionsTestProblem): Problem for extension test.
"""
problem.set_up()
backpack_res = BackpackExtensions(problem).sgs_extension_hook()
autograd_res = AutogradExtensions(problem).sgs()
check_sizes_and_values(autograd_res, backpack_res)
problem.tear_down()
def test_sum_grad_squared(problem):
"""Test sum of square of individual gradients
Args:
problem (ExtensionsTestProblem): Problem for extension test.
"""
problem.set_up()
backpack_res = BackpackExtensions(problem).sgs()
autograd_res = AutogradExtensions(problem).sgs()
check_sizes_and_values(autograd_res, backpack_res)
problem.tear_down()
def test_diag_ggn_batch(problem):
"""Test the individual diagonal of Generalized Gauss-Newton/Fisher
Args:
problem (ExtensionsTestProblem): Problem for extension test.
"""
problem.set_up()
backpack_res = BackpackExtensions(problem).diag_ggn_exact_batch()
autograd_res = AutogradExtensions(problem).diag_ggn_batch()
check_sizes_and_values(autograd_res, backpack_res)
problem.tear_down()
def test_batch_grad(problem):
"""Test individual gradients
Args:
problem (ExtensionsTestProblem): Problem for extension test.
"""
problem.set_up()
backpack_res = BackpackExtensions(problem).batch_grad()
autograd_res = AutogradExtensions(problem).batch_grad()
check_sizes_and_values(autograd_res, backpack_res)
problem.tear_down()
def test_diag_h_batch(problem):
"""Test Diagonal of Hessian
Args:
problem (ExtensionsTestProblem): Problem for extension test.
"""
problem.set_up()
backpack_res = BackpackExtensions(problem).diag_h_batch()
autograd_res = AutogradExtensions(problem).diag_h_batch()
check_sizes_and_values(autograd_res, backpack_res)
problem.tear_down()
def test_variance(problem: ExtensionsTestProblem) -> None:
"""Test variance of individual gradients.
Args:
problem: Test case.
"""
problem.set_up()
backpack_res = BackpackExtensions(problem).variance()
autograd_res = AutogradExtensions(problem).variance()
rtol = 5e-5
check_sizes_and_values(autograd_res, backpack_res, rtol=rtol)
problem.tear_down()
def test_diag_ggn(problem, request):
"""Test the diagonal of generalized Gauss-Newton.
Args:
problem (ExtensionsTestProblem): Problem for extension test.
request: problem request
"""
skip_adaptive_avg_pool3d_cuda(request)
problem.set_up()
backpack_res = BackpackExtensions(problem).diag_ggn()
autograd_res = AutogradExtensions(problem).diag_ggn()
check_sizes_and_values(autograd_res, backpack_res)
problem.tear_down()
def test_batch_grad(
problem: ExtensionsTestProblem, subsampling: Union[List[int], None]
) -> None:
"""Test individual gradients.
Args:
problem: Test case.
subsampling: Indices of active samples.
"""
skip_if_subsampling_conflict(problem, subsampling)
backpack_res = BackpackExtensions(problem).batch_grad(subsampling)
autograd_res = AutogradExtensions(problem).batch_grad(subsampling)
check_sizes_and_values(autograd_res, backpack_res)
def test_ggn_exact(
problem: ExtensionsTestProblem, subsampling: Union[List[int], None]
) -> None:
"""Compare exact GGN from BackPACK's matrix square root with autograd.
Args:
problem: Test case with small network whose GGN can be evaluated.
subsampling: Indices of active samples. ``None`` uses the full mini-batch.
"""
skip_large_parameters(problem)
skip_subsampling_conflict(problem, subsampling)
autograd_res = AutogradExtensions(problem).ggn(subsampling=subsampling)
backpack_res = BackpackExtensions(problem).ggn(subsampling=subsampling)
check_sizes_and_values(autograd_res, backpack_res)
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()