使用Captum搭配nn.Embedding时出现RuntimeError问题求助

RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.

• 核心问题：IntegratedGradients属于梯度类归因方法，要求输入张量支持梯度回传。你当前传入的输入是离散的token索引，属于整数类型，经过Embedding层时梯度只能回传到Embedding的参数，无法回传到输入的索引张量上，因此计算图找不到和输入张量关联的导数路径触发报错。
• 次要问题：你构造的输入indexes形状为(128,)，缺少batch维度，不符合模型前向传播的输入格式要求。

处理Embedding类输入的梯度归因，官方推荐使用LayerIntegratedGradients直接对Embedding层输出做归因，无需修改原有模型结构，修改后可正常运行的代码如下：

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
# 替换为LayerIntegratedGradients
from captum.attr import LayerIntegratedGradients

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

vocab_size = 1024
embedding_dim = 32
seq_len = 128
num_classes = 5
hidden_dim = 256

class predictor(nn.Module):
    def __init__(self):
        super().__init__()
        self.seq_len = seq_len
        self.num_classes = num_classes
        self.hidden_dim = hidden_dim 
        self.vocab_size, self.embedding_dim = vocab_size, embedding_dim

        self.embedding = nn.Embedding(self.vocab_size, self.embedding_dim)
        self.linear = nn.Linear(self.seq_len*self.embedding_dim, self.num_classes)

    def forward(self, x):
        x = self.embedding(x.long())
        x = x.reshape(-1, self.seq_len*self.embedding_dim)
        x = F.relu(self.linear(x))
        return x

class wrapper_predictor(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    def forward(self, x):
        x = self.model(x)
        x = F.softmax(x, dim=1)
        return x
    
# 给输入增加batch维度，形状变为(1, 128)
indexes = torch.Tensor(np.random.randint(0, vocab_size, (1, seq_len))).to(device)

model = predictor().to(device)
wrapper_model = wrapper_predictor(model).to(device)

# 初始化时指定要归因的层为Embedding层
lig = LayerIntegratedGradients(wrapper_model, model.embedding)
# 调用attribute方法，n_steps控制积分采样步数，步数越大结果越准确
attributions, delta = lig.attribute(inputs=indexes, target=0, n_steps=50, return_convergence_delta=True)

# 最终得到的attributions形状为(1, 128, 32)，对最后一维求和即可得到每个token的总贡献值
token_attributions = attributions.sum(dim=-1).squeeze(0)

内容的提问来源于stack exchange，提问作者OneQ