PyTorch中Max Pooling层为何存储输入张量？

我构建了一个包含卷积层和两个MaxPool2d层的简单模型：

class simple_model(nn.Module):
     def __init__(self):
         super(simple_model, self).__init__()
         self.maxpool2D = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
         self.conv1 = nn.Conv2d(3, 20, (5, 5))

     def forward(self, x):
         x = self.maxpool2D(self.maxpool2D(self.conv1(x)))

         return x

通过梯度钩子观察前向传播保存的张量，发现每个MaxPool2d层不仅保存了最大值索引的int64张量，还保存了输入的float32张量，且这些输入张量在反向传播中被使用。按我的理解，MaxPool反向传播只需要最大值索引就能完成梯度回传，为什么PyTorch还要保存输入张量？

这是因为PyTorch的MaxPool2d实现出于鲁棒性、兼容性和实现简便性的考虑，默认保留了输入张量的存储，核心原因包括：

• 处理数值异常场景：当输入张量中存在NaN或Inf时，仅靠最大值索引无法正确计算梯度——需要原始输入来判断哪些位置的梯度需要被置零或特殊处理，避免反向传播过程中出现数值崩溃。
• 兼容动态图灵活性：PyTorch的动态图特性允许运行时修改计算逻辑，保存输入张量能让反向传播过程更鲁棒，预留了对特殊自定义反向逻辑的兼容性。
• 历史实现与向后兼容：早期MaxPool反向传播实现依赖输入张量验证索引有效性，虽然后续可以仅靠索引完成计算，但为了不破坏旧代码的兼容性，没有彻底移除输入张量的保存逻辑。
• 算子实现的权衡：部分后端（如CUDA）的融合算子实现中，保存输入张量能简化前向-反向的融合逻辑，在内存开销和计算效率之间做了平衡。

如果需要严格控制内存，你可以通过自定义层，启用return_indices=True并仅保存索引来规避输入张量的存储：

class CustomMaxPool(nn.Module):
    def __init__(self):
        super().__init__()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0, return_indices=True)
    
    def forward(self, x):
        out, indices = self.pool(x)
        self.save_for_backward(indices)  # 仅保存索引，不存储输入x
        return out
    
    def backward(self, grad_output):
        indices, = self.saved_tensors
        # 仅用索引计算梯度，无需输入张量
        grad_input = torch.nn.functional.max_unpool2d(grad_output, indices, kernel_size=2, stride=2, padding=0)
        return grad_input

简单来说，虽然理论上仅靠索引就能完成MaxPool的反向传播，但PyTorch为了框架的鲁棒性和兼容性，默认选择保存输入张量。

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