def test_interface_sum_grad_squared():
interface_test(new_ext.SumGradSquared())
model = Sequential(Flatten(), Linear(784, 10),)
lossfunc = CrossEntropyLoss()
model = extend(model)
lossfunc = extend(lossfunc)
# %%
# We can now evaluate the loss and do a backward pass with Backpack
# -----------------------------------------------------------------
loss = lossfunc(model(X), y)
with backpack(
extensions.BatchGrad(),
extensions.Variance(),
extensions.SumGradSquared(),
extensions.BatchL2Grad(),
extensions.DiagGGNMC(mc_samples=1),
extensions.DiagGGNExact(),
extensions.DiagHessian(),
extensions.KFAC(mc_samples=1),
extensions.KFLR(),
extensions.KFRA(),
):
loss.backward()
# %%
# And here are the results
# -----------------------------------------------------------------
for name, param in model.named_parameters():
def test_interface_sum_grad_squared_conv():
interface_test(new_ext.SumGradSquared(), use_conv=True)
model, lossfunc, make_loader = utils.data_prep_cifar10()
backpack_secondorder_extension_classes = {
# "KFRA": extensions.KFRA,
"KFLR": extensions.KFLR,
"KFAC": extensions.KFAC,
"DiagGGNExact": extensions.DiagGGNExact,
"DiagGGNMC": extensions.DiagGGNMC,
"DiagH": extensions.DiagHessian,
}
backpack_extensions = {
"Var": extensions.Variance(),
"BatchGrad": extensions.BatchGrad(),
"BatchL2": extensions.BatchL2Grad(),
"SecondMoment": extensions.SumGradSquared(),
}
for name, ext_class in backpack_secondorder_extension_classes.items():
ext_class.add_module_extension(flatten, DiagGGNFlatten())
backpack_extensions[name] = ext_class()
combined_parameters = []
combined_names = []
for n in utils.Ns:
for name, ext in backpack_extensions.items():
combined_parameters.append((n, ext))
combined_names.append(str(n) + "-" + name)
@pytest.mark.parametrize("N, ext", combined_parameters, ids=combined_names)
def sgs(self):
with backpack(new_ext.SumGradSquared()):
_, _, loss = self.problem.forward_pass()
loss.backward()
sgs = [p.sum_grad_squared for p in self.problem.model.parameters()]
return sgs
def sgs(self):
with backpack(new_ext.SumGradSquared()):
self.loss().backward()
sgs = [p.sum_grad_squared for p in self.model.parameters()]
return sgs
"""
from torch.nn import CrossEntropyLoss, Flatten, Linear, Sequential
from backpack import backpack, extend, extensions
from backpack.utils.examples import load_mnist_data
B = 4
X, y = load_mnist_data(B)
print("# Gradient with PyTorch, gradient 2nd moment with BackPACK | B =", B)
model = Sequential(
Flatten(),
Linear(784, 10),
)
lossfunc = CrossEntropyLoss()
model = extend(model)
lossfunc = extend(lossfunc)
loss = lossfunc(model(X), y)
with backpack(extensions.SumGradSquared()):
loss.backward()
for name, param in model.named_parameters():
print(name)
print(".grad.shape: ", param.grad.shape)
print(".sum_grad_squared.shape: ", param.sum_grad_squared.shape)
def sgs(self) -> List[Tensor]: # noqa:D102
with backpack(new_ext.SumGradSquared()):
_, _, loss = self.problem.forward_pass()
loss.backward()
return self.problem.collect_data("sum_grad_squared")