IndexError: index 1 is out of bounds for dimension 1 with size 1 at DeepLift on custom VGG16

[martynanna]

❓ Questions and Help

Hey all, I'm trying to apply DeepLift on a custom VGG16, but getting an IndexError. Does somebody have an Idea how to fix it?

Reproduce

class CustomVGGNet(nn.Module):
    def __init__(self, modelname, nb_classes, dropout = 0.5, freeze = False):
        super().__init__()
        self.modelName = modelname
        self.dropout = dropout
        self.nb_classes = nb_classes
        self.model = {'vgg16': models.vgg16(pretrained = True),}
    
        self.pretrained_model = self.model[self.modelName]
        self.pooling_layer = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Sequential(
                           nn.Dropout(self.dropout),
                           nn.Linear(512, self.nb_classes))
        
    def forward(self, x):
        x = torch.squeeze(x, dim=0) 
        features = self.pretrained_model.features(x)
        pooled_features = self.pooling_layer(features)
        pooled_features = pooled_features.view(pooled_features.size(0), -1)
        flattened_features = torch.max(pooled_features, 0, keepdim=True)[0]
        output = self.classifier(flattened_features)
        return output


model = CustomVGGNet('vgg16',6)
model.eval()
model.to(torch.device('cuda')

dl = DeepLift(model)
image_tensor = torch.rand(28,3,158,158)
attr_dl = dl.attribute(image_tensor.cuda(), target=0)

Expected output

No errors expected (also no output on the console)

Actual output

The error:

Traceback (most recent call last):
  File "LRP.py", line 64, in <module>
    attr_dl = dl.attribute(image_tensor.cuda(), target = 0)
  File "/home/anna/.local/lib/python3.8/site-packages/captum/log/__init__.py", line 35, in wrapper
    return func(*args, **kwargs)
  File "/home/anna/.local/lib/python3.8/site-packages/captum/attr/_core/deep_lift.py", line 362, in attribute
    gradients = self.gradient_func(wrapped_forward_func, inputs)
  File "/home/anna/.local/lib/python3.8/site-packages/captum/_utils/gradient.py", line 112, in compute_gradients
    outputs = _run_forward(forward_fn, inputs, target_ind, additional_forward_args)
  File "/home/anna/.local/lib/python3.8/site-packages/captum/_utils/common.py", line 448, in _run_forward
    output = forward_func()
  File "/home/anna/.local/lib/python3.8/site-packages/captum/attr/_core/deep_lift.py", line 405, in forward_fn
    torch.cat((model_out[:, 0], model_out[:, 1])), target
IndexError: index 1 is out of bounds for dimension 1 with size 1

[swertz]

I'm running into the same issue. Why does DeepLift assume that the models outputs a 2D tensor?