Error when using the laplace package, a backpack module is raising issues

[MrHuff]

Hi!

Hope this message finds you well.

I am running laplace on a macbook (arm chips) using Pytorch 1.12 and keep getting this error:

AssertionError: BackPACK extension expects a backpropagation quantity but it is None. Module: Linear(in_features=128, out_features=1, bias=True), Extension: <backpack.extensions.secondorder.hbp.KFLR object at 0x169312c20>.

I suspect it's related to the hooks procedure that this pull request is patching.

Just wanted to clarify that downgrading to pytorch 1.8.x is likely to resolve this.

Thanks so much!

Best, Robert

[f-dangel]

Hi Robert, thanks for posting your problem!

what you're seeing is an error from BackPACK's KFLR extension, raised by (I assume due to its dimensions) the last layer of your network. Which loss function are you using? It seems that the information that BackPACK backpropagates for KFLR doesn't make it from the loss function to the last layer.

Do you have a short snippet to reproduce the issue?

Best, Felix

[MrHuff]

Hi Felix,

Thank you so much for your prompt response! Here is a minimal example to reproduce the error:

from sklearn.datasets import load_breast_cancer
import torch
from torch.utils.data import Dataset, DataLoader
from torch import nn
from torch.nn import functional as F
from laplace import Laplace

class Net(nn.Module):
    def __init__(self,input_shape):
        super(Net,self).__init__()
        self.fc1 = nn.Linear(input_shape,32)
        self.fc2 = nn.Linear(32,64)
        self.fc3 = nn.Linear(64,2)
    def forward(self,x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.softmax(self.fc3(x),dim=-1)
        return x

class dataset(Dataset):
    def __init__(self,x,y):
        self.x = torch.tensor(x,dtype=torch.float32)
        self.y = torch.tensor(y).long()
        self.length = self.x.shape[0]

    def __getitem__(self,idx):
        return self.x[idx],self.y[idx]
    def __len__(self):
        return self.length
data = load_breast_cancer()
x = data['data']
y = data['target']
print("shape of x: {}\nshape of y: {}".format(x.shape,y.shape))

trainset = dataset(x,y)
#DataLoader
trainloader = DataLoader(trainset,batch_size=64,shuffle=False)
model = Net(input_shape=x.shape[1])
la = Laplace(model,'classification')
la.fit(trainloader)
la.optimize_prior_precision()

Code run on a MacBook Pro (14-inch, 2021) with Apple M1 Pro chips and Pytorch 1.12.1 in cpu mode. The dependencies are fairly standard so it should run out of the box with almost no issues.

Thanks!

Best, Robert

[f-dangel]

Hi Robert,

the following seems to work (explanations below):

from laplace import Laplace
from sklearn.datasets import load_breast_cancer
from torch import from_numpy, nn
from torch.utils.data import DataLoader, TensorDataset

def net(D_in):
    return nn.Sequential(
        nn.Linear(D_in, 32),
        nn.ReLU(),
        nn.Linear(32, 64),
        nn.ReLU(),
        nn.Linear(64, 2),
    )

data = load_breast_cancer()
x, y = from_numpy(data["data"]).float(), from_numpy(data["target"]).long()

trainset = TensorDataset(x, y)
trainloader = DataLoader(trainset, batch_size=64, shuffle=False)

model = net(x.shape[1])
la = Laplace(model, "classification")
la.fit(trainloader)
la.optimize_prior_precision()

• In order for your model to be compatible with BackPACK, it has to consist of nn.Modules that are supported (more details here).
• I believe that when you set "classification" in the Laplace approximation, it will use nn.CrossEntropyLoss under the hood, which already performs a softmax. So I removed your softmax from the network. Maybe double-check this with the laplace maintainers if you run into problems with this.

Best, Felix

[MrHuff]

Thanks so much for the help Felix!

I'll try this solution and get back!

Best, Robert

[f-dangel]

You're welcome!

Feel free to close this issue if you consider it resolved.

[MrHuff]

Managed to completely fix the issues now. Thanks so much for the help Felix!