How to use BackPACK

The use BackPACK with your setup, you first need to backpack.extend() the model and the loss function and register the extension you want to use with backpack.backpack() before calling the backward() function

Extending the model and loss function

import torch

model = torch.nn.Sequential(
        torch.nn.Linear(764, 64),
        torch.nn.ReLU(),
        torch.nn.Linear(64, 10)
)
lossfunc = torch.nn.CrossEntropyLoss()

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

See Supported models for the list of supported layers.

backpack.extend(module, debug=False)

Extends the module to make it backPACK-ready.

module: torch.nn.Module

The module to extend

debug: Bool, optional (default: False)

If true, will print debug messages during the extension and backward.

Calling the extension

from backpack import backpack
from backpack.extensions import KFAC
from utils import load_data

X, y = load_data()

loss = lossfunc(model(X), y)

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

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

See Extensions for the list of available extensions and how to access the quantities.

backpack.backpack(*args)

Activates the BackPACK extensions passed as arguments for the backward calls in the current with block.