Pytorch 将Captum与Pytork Lightning一起使用？

所以我试着用Captum和PyTorch Lightning。我在将模块传递给Captum时遇到了问题，因为它似乎对张量进行了奇怪的重塑。 例如，在下面的最小示例中，lightning代码工作简单且良好。 但是当我使用积分梯度与“n_步>=1”我得到一个问题。 LighningModule的代码并没有那么重要，我想说的是，我更想知道最底层的代码行

有人知道如何解决这个问题吗

from captum.attr import IntegratedGradients
from torch import nn, optim, rand, sum as tsum, reshape, device
import torch.nn.functional as F
from pytorch_lightning import seed_everything, LightningModule, Trainer
from torch.utils.data import DataLoader, Dataset

SAMPLE_DIM = 3


class CustomDataset(Dataset):
    def __init__(self, samples=42):
        self.dataset = rand(samples, SAMPLE_DIM).cuda().float() * 2 - 1

    def __getitem__(self, index):
        return (self.dataset[index], (tsum(self.dataset[index]) > 0).cuda().float())

    def __len__(self):
        return self.dataset.size()[0]


class OurModel(LightningModule):
    def __init__(self):
        super(OurModel, self).__init__()
        # Network layers
        self.linear = nn.Linear(SAMPLE_DIM, 2048)
        self.linear2 = nn.Linear(2048, 1)
        self.output = nn.Sigmoid()
        # Hyper-parameters, that we will auto-tune using lightning!
        self.lr = 0.001
        self.batch_size = 512

    def forward(self, x):
        x = self.linear(x)
        x = self.linear2(x)
        output = self.output(x)
        return reshape(output, (-1,))

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.lr)

    def train_dataloader(self):
        loader = DataLoader(CustomDataset(samples=1000), batch_size=self.batch_size, shuffle=True)
        return loader

    def training_step(self, batch, batch_nb):
        x, y = batch
        loss = F.binary_cross_entropy(self(x), y)
        return {'loss': loss, 'log': {'train_loss': loss}}


if __name__ == '__main__':
    seed_everything(42)
    device = device("cuda")
    model = OurModel().to(device)
    trainer = Trainer(max_epochs=2, min_epochs=1, auto_lr_find=False,
                      progress_bar_refresh_rate=10)
    trainer.fit(model)
    # ok Now the Problem
    test_input = CustomDataset(samples=1).__getitem__(0)[0].requires_grad_()
    ig = IntegratedGradients(model)
    attr, delta = ig.attribute(test_input, target=1, return_convergence_delta=True)

---

解决方案是包装forward函数。确保进入模式.foward（）的形状正确

# Solution is this wrapper function
def modified_f(in_vec):
    # Shape here is wrong
    print("IN:", in_vec.size())
    x = torch.reshape(in_vec, (int(in_vec.size()[0]/SAMPLE_DIM), SAMPLE_DIM))
    print("x:", x.size())

    res = model.forward(x)
    print("res:", res.size())
    res = torch.reshape(res, (res.size()[0], 1))
    print("res2:", res.size())

    return res


ig = IntegratedGradients(modified_f)
attr, delta = ig.attribute(test_input, return_convergence_delta=True, n_steps=STEP_AMOUNT)