BackPACK is a library built on top of PyTorch to make it easy to extract more information from a backward pass. Some of the things you can compute:

- the gradient with PyTorch
- an estimate of the Variance
- the Gauss-Newton Diagonal
- the Gauss-Newton KFAC

    """
    Compute the gradient with Pytorch

    """
    from torch.nn import CrossEntropyLoss, Linear
    from utils import load_mnist_data

    X, y = load_mnist_data()
    model = Linear(784, 10)
    lossfunc = CrossEntropyLoss()
    loss = lossfunc(model(X), y)

    loss.backward()

    for param in model.parameters():
        print(param.grad)

    """
    Compute the gradient with Pytorch
    and the variance with BackPACK
    """
    from torch.nn import CrossEntropyLoss, Linear
    from utils import load_mnist_data
    from backpack import extend, backpack, Variance

    X, y = load_mnist_data()
    model = extend(Linear(784, 10))
    lossfunc = extend(CrossEntropyLoss())
    loss = lossfunc(model(X), y)

    with backpack(Variance()):
        loss.backward()

    for param in model.parameters():
        print(param.grad)
        print(param.variance)

    """
    Compute the gradient with Pytorch
    and the second moment with BackPACK
    """
    from torch.nn import CrossEntropyLoss, Linear
    from utils import load_mnist_data
    from backpack import extend, backpack, SumGradSquared

    X, y = load_mnist_data()
    model = extend(Linear(784, 10))
    lossfunc = extend(CrossEntropyLoss())
    loss = lossfunc(model(X), y)

    with backpack(SumGradSquared()):
        loss.backward()

    for param in model.parameters():
        print(param.grad)
        print(param.sum_grad_squared)

    """
    Compute the gradient with Pytorch
    and the diagonal of the Gauss-Newton with BackPACK
    """
    from torch.nn import CrossEntropyLoss, Linear
    from utils import load_mnist_data
    from backpack import extend, backpack, DiagGGNExact

    X, y = load_mnist_data()
    model = extend(Linear(784, 10))
    lossfunc = extend(CrossEntropyLoss())
    loss = lossfunc(model(X), y)

    with backpack(DiagGGNExact()):
        loss.backward()

    for param in model.parameters():
        print(param.grad)
        print(param.diag_ggn_exact)

    """
    Compute the gradient with Pytorch
    and KFAC with BackPACK
    """
    from torch.nn import CrossEntropyLoss, Linear
    from utils import load_mnist_data
    from backpack import extend, backpack, KFAC

    X, y = load_mnist_data()
    model = extend(Linear(784, 10))
    lossfunc = extend(CrossEntropyLoss())
    loss = lossfunc(model(X), y)

    with backpack(KFAC()):
        loss.backward()

    for param in model.parameters():
        print(param.grad)
        print(param.kfac)

Install with

    pip install backpack-for-pytorch

If you use BackPACK in your research, please cite download bibtex

    @inproceedings{dangel2020backpack,
        title = {BackPACK: Packing more into Backprop},
        author = {Felix Dangel and Frederik Kunstner and Philipp Hennig},
        booktitle = {International Conference on Learning Representations},
        year = {2020},
        url = {https://openreview.net/forum?id=BJlrF24twB}
    }