Pytorch - 香草梯度可视化

Pytorch - vanilla Gradient Visualization

我使用 PyTorch 在 MNIST 上训练了一个神经网络:

class MnistCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d( 1, 16, 3, stride = 1, padding = 2)
        self.pool1 = nn.MaxPool2d(kernel_size = 2)
        self.conv2 = nn.Conv2d(16, 32, 3, stride = 1, padding = 2)
        self.pool2 = nn.MaxPool2d(kernel_size = 2)
        self.dropout = nn.Dropout(0.5)
        self.lin = nn.Linear(32 * 8 * 8, 10)
    
    def forward(self, x):
        # conv block
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        # conv block
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        # dense block
        x = x.view(x.size(0), -1)
        x = self.dropout(x)
        return self.lin(x)

我想在我的模型上实现香草梯度可视化(参见下面的参考资料)。

Simonyan, K., Vedaldi, A., Zisserman, A.
Deep inside convolutional networks: Visualising image classification models and saliency maps.
arXiv preprint arXiv:1312.6034 (2013)

问题:如何在PyTorch中实现这个方法?

如果我理解正确,vanilla 梯度可视化 在于计算模型损失的偏导数 w.r.t 输入图像中的所有像素。所以简而言之,我需要调整我的 self.conv1 层,以便它计算输入像素的梯度而不是权重的梯度。

如有错误请指正

你不需要改变你的转换层的任何东西。每层计算梯度 both w.r.t。参数(用于更新) w.r.t。输入(对于链式法则的“下游”梯度)。因此,您只需将输入图像的 x 梯度 属性 设置为 true:

x, y = ... # get one image from MNIST
x.requires_grad_(True)  # indicate to pytorch that you would like to look at these gradients
pred = model(x)
loss = criterion(pred, y)
loss.backward()  # propagate gradients
x.grad  # <- here you should have the gradients of the loss w.r.t pixels