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():
print(name)
print(".grad.shape: ", param.grad.shape)
print(".grad_batch.shape: ", param.grad_batch.shape)
print(".variance.shape: ", param.variance.shape)
def test_interface_diag_h():
interface_test(new_ext.DiagHessian())
def diag_h(self):
with backpack(new_ext.DiagHessian()):
_, _, loss = self.problem.forward_pass()
loss.backward()
diag_h = [p.diag_h for p in self.problem.model.parameters()]
return diag_h
def diag_h(self):
with backpack(new_ext.DiagHessian()):
self.loss().backward()
diag_h = [p.diag_h for p in self.model.parameters()]
return diag_h
"""
Compute the gradient with PyTorch and the Hessian diagonal with BackPACK.
"""
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, Hessian diagonal 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.DiagHessian()):
loss.backward()
for name, param in model.named_parameters():
print(name)
print(".grad.shape: ", param.grad.shape)
print(".diag_h.shape: ", param.diag_h.shape)
def diag_h(self) -> List[Tensor]: # noqa:D102
with backpack(new_ext.DiagHessian()):
_, _, loss = self.problem.forward_pass()
loss.backward()
return self.problem.collect_data("diag_h")
plotter = CockpitPlotter()
# Main training loop
max_steps, global_step = 50, 0
for inputs, labels in iter(fmnist_data):
opt.zero_grad()
# forward pass
outputs = model(inputs)
loss = loss_fn(outputs, labels)
losses = individual_loss_fn(outputs, labels)
# backward pass
with cockpit(
global_step,
extensions.DiagHessian(
), # Other BackPACK quantities can be computed as well
info={
"batch_size": inputs.shape[0],
"individual_losses": losses,
"loss": loss,
"optimizer": opt,
},
):
loss.backward(create_graph=cockpit.create_graph(global_step))
# optimizer step
opt.step()
global_step += 1
print(f"Step: {global_step:5d} | Loss: {loss.item():.4f}")