Note

Click here to download the full example code

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 torch import allclose, rand
from torch.nn import CrossEntropyLoss, Flatten, Linear, Sequential

from backpack import backpack, extend
from backpack.extensions import (
    GGNMP,
    HMP,
    KFAC,
    KFLR,
    KFRA,
    PCHMP,
    BatchGrad,
    BatchL2Grad,
    DiagGGNExact,
    DiagGGNMC,
    DiagHessian,
    SumGradSquared,
    Variance,
)
from backpack.utils.examples import load_one_batch_mnist

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...
/home/docs/checkouts/readthedocs.org/user_builds/backpack/envs/development/lib/python3.7/site-packages/torchvision/datasets/mnist.py:480: 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. (Triggered internally at  /pytorch/torch/csrc/utils/tensor_numpy.cpp:141.)
  return torch.from_numpy(parsed.astype(m[2], copy=False)).view(*s)
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:

/home/docs/checkouts/readthedocs.org/user_builds/backpack/envs/development/lib/python3.7/site-packages/torch/autograd/__init__.py:132: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  /pytorch/c10/cuda/CUDAFunctions.cpp:100.)
  allow_unreachable=True)  # allow_unreachable flag
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])

Block-diagonal curvature products¶

Curvature-matrix product (MP) extensions provide functions that multiply with the block diagonal of different curvature matrices, such as

the Hessian (HMP)
the generalized Gauss-Newton (GGNMP)
the positive-curvature Hessian (PCHMP)

loss = lossfunc(model(X), y)

with backpack(
    HMP(),
    GGNMP(),
    PCHMP(savefield="pchmp_clip", modify="clip"),
    PCHMP(savefield="pchmp_abs", modify="abs"),
):
    loss.backward()

Multiply a random vector with curvature blocks.

V = 1

for name, param in model.named_parameters():
    vec = rand(V, *param.shape)
    print(name)
    print(".grad.shape:             ", param.grad.shape)
    print("vec.shape:               ", vec.shape)
    print(".hmp(vec).shape:         ", param.hmp(vec).shape)
    print(".ggnmp(vec).shape:       ", param.ggnmp(vec).shape)
    print(".pchmp_clip(vec).shape:  ", param.pchmp_clip(vec).shape)
    print(".pchmp_abs(vec).shape:   ", param.pchmp_abs(vec).shape)

Out:

1.weight
.grad.shape:              torch.Size([10, 784])
vec.shape:                torch.Size([1, 10, 784])
.hmp(vec).shape:          torch.Size([1, 10, 784])
.ggnmp(vec).shape:        torch.Size([1, 10, 784])
.pchmp_clip(vec).shape:   torch.Size([1, 10, 784])
.pchmp_abs(vec).shape:    torch.Size([1, 10, 784])
1.bias
.grad.shape:              torch.Size([10])
vec.shape:                torch.Size([1, 10])
.hmp(vec).shape:          torch.Size([1, 10])
.ggnmp(vec).shape:        torch.Size([1, 10])
.pchmp_clip(vec).shape:   torch.Size([1, 10])
.pchmp_abs(vec).shape:    torch.Size([1, 10])

Multiply a collection of three vectors (a matrix) with curvature blocks.

V = 3

for name, param in model.named_parameters():
    vec = rand(V, *param.shape)
    print(name)
    print(".grad.shape:             ", param.grad.shape)
    print("vec.shape:               ", vec.shape)
    print(".hmp(vec).shape:         ", param.hmp(vec).shape)
    print(".ggnmp(vec).shape:       ", param.ggnmp(vec).shape)
    print(".pchmp_clip(vec).shape:  ", param.pchmp_clip(vec).shape)
    print(".pchmp_abs(vec).shape:   ", param.pchmp_abs(vec).shape)

Out:

1.weight
.grad.shape:              torch.Size([10, 784])
vec.shape:                torch.Size([3, 10, 784])
.hmp(vec).shape:          torch.Size([3, 10, 784])
.ggnmp(vec).shape:        torch.Size([3, 10, 784])
.pchmp_clip(vec).shape:   torch.Size([3, 10, 784])
.pchmp_abs(vec).shape:    torch.Size([3, 10, 784])
1.bias
.grad.shape:              torch.Size([10])
vec.shape:                torch.Size([3, 10])
.hmp(vec).shape:          torch.Size([3, 10])
.ggnmp(vec).shape:        torch.Size([3, 10])
.pchmp_clip(vec).shape:   torch.Size([3, 10])
.pchmp_abs(vec).shape:    torch.Size([3, 10])

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

Gallery generated by Sphinx-Gallery