将 Captum 与 Pytorch Lightning 一起使用？

Posted 2023-03-27

【中文标题】将 Captum 与 Pytorch Lightning 一起使用？

【英文标题】：Using Captum with Pytorch Lightning?

【发布时间】：2021-01-23 08:07:35

【问题描述】：

所以我尝试将 Captum 与 PyTorch Lightning 一起使用。将模块传递给 Captum 时遇到问题，因为它似乎对张量进行了奇怪的重塑。 例如，在下面的最小示例中，闪电代码运行起来既简单又好。 但是当我将 IntegratedGradient 与“n_step>=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)

【参考方案1】：

解决方案是包装 forward 函数。确保进入 mode.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)