1. 使用自定义函数的双重反向传播
有时候,在反向计算图中运行两次反向传播是有用的,例如计算高阶梯度。然而,支持双重反向传播需要对自动求导(autograd)有一定的理解,并且需要小心处理。支持单次反向传播的函数不一定能够支持双重反向传播。在本教程中,我们将展示如何编写一个支持双重反向传播的自定义自动求导函数,并指出一些需要注意的事项。
在编写一个支持两次反向传播的自定义自动求导函数时,了解自定义函数中的操作何时被自动求导记录、何时不被记录,以及最重要的是,save_for_backward 如何与这些机制配合工作,是非常关键的。
自定义函数以两种方式隐式影响梯度模式:
-
在前向传播期间,自动求导不会记录在前向函数中执行的任何操作的计算图。当前向传播完成时,自定义函数的反向函数将成为每个前向输出的 grad_fn。
-
在反向传播期间,如果指定了 create_graph,自动求导会记录用于计算反向传播的计算图。
接下来,为了理解 save_for_backward 如何与上述机制交互,我们可以通过几个示例来探讨。
1.1保存输入
考虑这个简单的平方函数。它保存了一个输入张量以便用于反向传播。双重反向传播会在 autograd 能够记录反向传播中的操作时自动