排查PyTorch中模型输出对神经元预激活的偏导数计算代码问题

我在研究分类场景下人工神经网络（ANN）的神经元重要性，采用的基线方法是计算正确类别的模型输出对目标神经元预激活的偏导数。我的目标是获取对应输入簇的重要神经元，因此对每个簇内样本的重要性得分取平均得到簇级重要性。但目前遇到两个问题：

• 每个神经元的平均偏导数数值极小
• 基于该得分移除“最关键”神经元时，模型性能下降幅度和随机移除无差异；反而用预激活值作为重要性得分的效果更好

其他研究者的实验显示，基于偏导数的方法在多数场景下效果合理，因此我怀疑代码存在bug。

当前代码实现如下（以清晰性为编写原则，未做性能优化）：

from abc import ABC, abstractmethod
import torch
from torch import nn
from torch.utils.data import DataLoader
import numpy as np

class AttributionMethod(ABC):
    def __init__(self, model: nn.Module, data_loader: dict[int, DataLoader]):
        self.model = model
        self.data_loader = data_loader
        if self.model is not None:
            self.model.eval()

    @abstractmethod
    def attribute(self, layer: nn.Module, latent_class: int) -> np.ndarray:
        pass


class GradientAttribution(AttributionMethod):
    def attribute(self, layer: nn.Module, latent_class: int) -> np.ndarray:
        # 计算指定层中对应簇的神经元重要性得分
        device = next(self.model.parameters()).device

        inputs, targets = next(iter(self.data_loader[latent_class]))
        inputs = inputs.to(device)
        targets = targets.to(device)
        inputs.requires_grad = True

        attributions = []

        def hook_fn(module, grad_input, grad_output):
            attributions.append(grad_output[0].detach().cpu())

        handle = layer.register_full_backward_hook(hook_fn)

        # 前向与反向传播
        for i in range(len(inputs)):
            input_i = inputs[i].unsqueeze(0)  # 添加batch维度
            output = self.model(input_i).squeeze()
            target_i = targets[i].item()

            output[target_i].backward()
            # 注释：认为无需清零梯度，因为只有模型参数梯度会累积

        handle.remove()

        # 拼接batch结果并按神经元取平均
        grads_tensor = torch.cat(attributions, dim=0)
        grads_per_neuron = grads_tensor.mean(dim=0).numpy()
        return grads_per_neuron

代码中的核心错误 #

1. 梯度未清零导致累积错误
   每次调用backward()时，PyTorch会自动累积模型参数和输入的梯度。你循环处理每个样本时没有清零梯度，导致后面样本的梯度是之前所有样本梯度的总和，最终计算的平均梯度完全偏离真实值。这是导致得分无区分度、和随机移除效果一致的核心原因。

2. 输入梯度的潜在污染
   原输入inputs设置了requires_grad=True，但循环中复用该输入的切片，每次反向后输入的梯度会保留并叠加，进一步干扰后续样本的梯度计算。

3. 钩子函数的理解偏差（非直接错误，但需明确）
   对于Linear层，register_full_backward_hook的grad_output参数就是模型正确类别输出对该层预激活的偏导数，你取grad_output[0]是正确的（因为Linear层只有一个输出张量，元组中仅含一个元素），但梯度累积错误导致这个值完全不可靠。

修正后的关键代码部分 #

# 前向与反向传播
for i in range(len(inputs)):
    # 为每个样本创建独立的输入张量并开启梯度
    input_i = inputs[i].unsqueeze(0).to(device).requires_grad_(True)
    output = self.model(input_i).squeeze()
    target_i = targets[i].item()

    # 必须清零模型参数梯度和输入梯度
    self.model.zero_grad()
    if input_i.grad is not None:
        input_i.grad.zero_()

    output[target_i].backward()

额外验证建议 #

• 先测试单个样本：取一个样本执行前向+反向，手动计算输出对预激活的梯度（比如对Linear层，预激活是wx+b，输出对它的梯度等于该层输出到最终分类层的梯度传递值），和钩子捕获的结果对比，验证梯度计算是否正确。
• 检查梯度的分布：如果修正后梯度仍然极小，可能是模型已收敛，输出对预激活的灵敏度本身就低，但此时得分应该仍有区分度，不会和随机移除效果一致。

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