Example using all extensions¶

Basic example showing how compute the gradient, and and other quantities with BackPACK, on a linear model for MNIST.

Let’s start by loading some dummy data and extending the model

from backpack.utils.examples import load_one_batch_mnist
from torch.nn import CrossEntropyLoss, Flatten, Linear, Sequential
from backpack import backpack, extend
from backpack.extensions import KFAC, KFLR, KFRA
from backpack.extensions import DiagGGNExact, DiagGGNMC, DiagHessian
from backpack.extensions import BatchGrad, SumGradSquared, Variance, BatchL2Grad

X, y = load_one_batch_mnist(batch_size=512)

model = Sequential(Flatten(), Linear(784, 10),)
lossfunc = CrossEntropyLoss()

model = extend(model)
lossfunc = extend(lossfunc)

Out:

Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./data/MNIST/raw/train-images-idx3-ubyte.gz

0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
2%
2%
3%
4%
5%
6%
7%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
1%
1%
2%
3%
4%
5%
6%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
0%
0%
1%
2%
3%
4%
5%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
4%
4%
5%
6%
7%
8%
9%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%
2%
3%
3%
4%
5%
6%
7%
8%
8%
9%
0%
1%Extracting ./data/MNIST/raw/train-images-idx3-ubyte.gz to ./data/MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./data/MNIST/raw/train-labels-idx1-ubyte.gz

0%
4%
7%
1%
5%Extracting ./data/MNIST/raw/train-labels-idx1-ubyte.gz to ./data/MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw/t10k-images-idx3-ubyte.gz

0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
8%
3%
8%
3%
8%
3%
8%
3%
8%
3%
8%
3%
8%
3%
8%
3%
8%
3%
8%
3%
8%
3%
8%
3%
8%
3%
8%
3%
8%
3%
8%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
2%
7%
1%
6%
1%
6%
1%
6%
1%
6%
1%
6%
1%
6%
1%
6%
1%
6%
1%
6%
1%
6%
1%
6%
1%
6%
1%
6%
1%
6%
1%
6%
1%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
5%
0%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%
9%
4%Extracting ./data/MNIST/raw/t10k-images-idx3-ubyte.gz to ./data/MNIST/raw
Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz

0%
4%Extracting ./data/MNIST/raw/t10k-labels-idx1-ubyte.gz to ./data/MNIST/raw
Processing...
/pytorch/torch/csrc/utils/tensor_numpy.cpp:141: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program.
Done!

First order extensions¶

Batch gradients

loss = lossfunc(model(X), y)
with backpack(BatchGrad()):
    loss.backward()

for name, param in model.named_parameters():
    print(name)
    print(".grad.shape:             ", param.grad.shape)
    print(".grad_batch.shape:       ", param.grad_batch.shape)

Out:

1.weight
.grad.shape:              torch.Size([10, 784])
.grad_batch.shape:        torch.Size([512, 10, 784])
1.bias
.grad.shape:              torch.Size([10])
.grad_batch.shape:        torch.Size([512, 10])

Variance

loss = lossfunc(model(X), y)
with backpack(Variance()):
    loss.backward()

for name, param in model.named_parameters():
    print(name)
    print(".grad.shape:             ", param.grad.shape)
    print(".variance.shape:         ", param.variance.shape)

Out:

1.weight
.grad.shape:              torch.Size([10, 784])
.variance.shape:          torch.Size([10, 784])
1.bias
.grad.shape:              torch.Size([10])
.variance.shape:          torch.Size([10])

Second moment/sum of gradients squared

loss = lossfunc(model(X), y)
with backpack(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)

Out:

1.weight
.grad.shape:              torch.Size([10, 784])
.sum_grad_squared.shape:  torch.Size([10, 784])
1.bias
.grad.shape:              torch.Size([10])
.sum_grad_squared.shape:  torch.Size([10])

L2 norm of individual gradients

loss = lossfunc(model(X), y)
with backpack(BatchL2Grad()):
    loss.backward()

for name, param in model.named_parameters():
    print(name)
    print(".grad.shape:             ", param.grad.shape)
    print(".batch_l2.shape:         ", param.batch_l2.shape)

Out:

1.weight
.grad.shape:              torch.Size([10, 784])
.batch_l2.shape:          torch.Size([512])
1.bias
.grad.shape:              torch.Size([10])
.batch_l2.shape:          torch.Size([512])

It’s also possible to ask for multiple quantities at once

loss = lossfunc(model(X), y)
with backpack(BatchGrad(), Variance(), SumGradSquared(), BatchL2Grad()):
    loss.backward()

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)
    print(".sum_grad_squared.shape: ", param.sum_grad_squared.shape)
    print(".batch_l2.shape:         ", param.batch_l2.shape)

Out:

1.weight
.grad.shape:              torch.Size([10, 784])
.grad_batch.shape:        torch.Size([512, 10, 784])
.variance.shape:          torch.Size([10, 784])
.sum_grad_squared.shape:  torch.Size([10, 784])
.batch_l2.shape:          torch.Size([512])
1.bias
.grad.shape:              torch.Size([10])
.grad_batch.shape:        torch.Size([512, 10])
.variance.shape:          torch.Size([10])
.sum_grad_squared.shape:  torch.Size([10])
.batch_l2.shape:          torch.Size([512])

Second order extensions¶

Diagonal of the Gauss-Newton and its Monte-Carlo approximation

loss = lossfunc(model(X), y)
with backpack(DiagGGNExact(), DiagGGNMC(mc_samples=1)):
    loss.backward()

for name, param in model.named_parameters():
    print(name)
    print(".grad.shape:             ", param.grad.shape)
    print(".diag_ggn_mc.shape:      ", param.diag_ggn_mc.shape)
    print(".diag_ggn_exact.shape:   ", param.diag_ggn_exact.shape)

Out:

1.weight
.grad.shape:              torch.Size([10, 784])
.diag_ggn_mc.shape:       torch.Size([10, 784])
.diag_ggn_exact.shape:    torch.Size([10, 784])
1.bias
.grad.shape:              torch.Size([10])
.diag_ggn_mc.shape:       torch.Size([10])
.diag_ggn_exact.shape:    torch.Size([10])

KFAC, KFRA and KFLR

loss = lossfunc(model(X), y)
with backpack(KFAC(mc_samples=1), KFLR(), KFRA()):
    loss.backward()

for name, param in model.named_parameters():
    print(name)
    print(".grad.shape:             ", param.grad.shape)
    print(".kfac (shapes):          ", [kfac.shape for kfac in param.kfac])
    print(".kflr (shapes):          ", [kflr.shape for kflr in param.kflr])
    print(".kfra (shapes):          ", [kfra.shape for kfra in param.kfra])

Out:

1.weight
.grad.shape:              torch.Size([10, 784])
.kfac (shapes):           [torch.Size([10, 10]), torch.Size([784, 784])]
.kflr (shapes):           [torch.Size([10, 10]), torch.Size([784, 784])]
.kfra (shapes):           [torch.Size([10, 10]), torch.Size([784, 784])]
1.bias
.grad.shape:              torch.Size([10])
.kfac (shapes):           [torch.Size([10, 10])]
.kflr (shapes):           [torch.Size([10, 10])]
.kfra (shapes):           [torch.Size([10, 10])]

Diagonal Hessian

loss = lossfunc(model(X), y)
with backpack(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)

Out:

1.weight
.grad.shape:              torch.Size([10, 784])
.diag_h.shape:            torch.Size([10, 784])
1.bias
.grad.shape:              torch.Size([10])
.diag_h.shape:            torch.Size([10])

Total running time of the script: ( 0 minutes 1.974 seconds)

Gallery generated by Sphinx-Gallery