如何使用PyTorch计算非叶节点变量的梯度示例？

本文共计537个文字，预计阅读时间需要3分钟。

在PyTorch中，通常只对叶节点进行梯度计算，即下图中d、e节点，而非叶节点如c、b节点则没有显式地保存其梯度。这是因为只有叶节点才需要更新，而非叶节点的梯度可以通过链式法则从叶节点反向传播得到（因为只有叶节点才需要更新）。

在pytorch中一般只对叶节点进行梯度计算，也就是下图中的d,e节点，而对非叶节点，也即是c,b节点则没有显式地去保留其中间计算过程中的梯度（因为一般来说只有叶节点才需要去更新），这样可以节省很大部分的显存，但是在调试过程中，有时候我们需要对中间变量梯度进行监控，以确保网络的有效性，这个时候我们需要打印出非叶节点的梯度，为了实现这个目的，我们可以通过两种手段进行。

注册hook函数

Tensor.register_hook[2] 可以注册一个反向梯度传导时的hook函数，这个hook函数将会在每次计算 关于该张量 的时候 被调用，经常用于调试的时候打印出非叶节点梯度。当然，通过这个手段，你也可以自定义某一层的梯度更新方法。