一枚NLPer小菜鸡

Pytorch Autograd 机制介绍

发表于 更新于

什么是 Autograd?

Autograd是反向自动微分系统。从概念上讲，autograd 记录一个图结构，记录执行操作时创建数据的所有操作，一个有向无环图，其叶子是输入张量，根是输出张量。通过从根到叶跟踪此图，可以使用链式法则自动计算梯度。

更为具体的来说，Autograd方法是构造一个图结构的函数传播形式。它在计算网络的前向传播过程中构造了一个梯度的函数计算图。当完成前向过程后，进行反向传播时，计算每一个节点的梯度。在每一轮次都是即时构建这个图结构，并不用提前对图结构进行编码。按照python的函数流程，所构即所得。

保存Tensor张量

一些操作需要在前向传递期间保存中间结果，以便执行反向传递。例如，函数 $x\mapsto x^{2}$ 保存输入 $x$ 以计算梯度。

定义自定义 Python 函数时，您可以使用 save_for_backward() 在前向传递期间保存张量，并在后向传递期间使用 saved_tensors 检索它们。

对于PyTorch定义的操作（例如 torch.pow()），张量会根据需要自动保存。您可以通过查找以前缀 _saved 开头的属性来探索某个grad_fn 保存了哪些张量。

1
2
3
4
x = torch.randn(5, requires_grad=True)
y = x.pow(2)
print(x.equal(y.grad_fn._saved_self))  # True
print(x is y.grad_fn._saved_self)  # True

在前面的代码中，y.grad_fn._saved_self 指的是与 x 相同的 Tensor 对象。但情况可能并非总是如此。例如：

1
2
3
4
x = torch.randn(5, requires_grad=True)
y = x.exp()
print(y.equal(y.grad_fn._saved_result))  # True
print(y is y.grad_fn._saved_result)  # False

在底层，为了防止引用循环，PyTorch 在保存时打包了张量，并将其解包到不同的张量中以供读取。在这里，您通过访问 y.grad_fn._saved_result 获得的张量是与 y 不同的张量对象（但它们仍然共享相同的存储）。一个张量是否会被打包到不同的张量对象中，取决于它是否是它自己的 grad_fn 的输出，这是一个可能会发生变化的实现细节，因此在操作时不应该依赖它。

局部禁用梯度计算

Python 有几种机制可以局部的禁用梯度计算：

要禁用整个代码块的梯度，有上下文管理器，如 no-grad 模式和 inference 模式。为了从梯度计算中更细粒度地排除子图，可以设置张量的 requires_grad 字段。
下面，除了讨论上述机制之外，我们还描述了 evaluate 模式nn.Module.eval()，这种方法实际上并不用于禁用梯度计算，但由于其名称，经常与这三种方法混淆。

设置 required_grad

requires_grad 是一个标志，除非包含在 nn.Parameter 中，否则默认为 false，它允许从梯度计算中细粒度地排除子图。它在向前和向后传递中都生效：

在前向传递期间，只有至少一个输入张量需要 grad 时，才会将操作记录在后向图中。在后向传递 (.backward()) 期间，只有 requires_grad=True 的叶张量才会将梯度累积到它们的 .grad 字段中。

重要的是要注意，即使每个张量都有这个标志，设置它只对叶张量有意义（没有 grad_fn 的张量，例如，nn.Module 的参数）。非叶张量（具有 grad_fn 的张量）是具有与之关联的后向图的张量。因此，需要它们的梯度作为中间结果来计算需要 grad 的叶张量的梯度。从这个定义中，很明显所有非叶张量都会自动具有require_grad=True。

设置require_grad字段是有效的控制在你所用的模型中哪一部分的梯度需要被计算。比如说，当你使用预训练模型Bert类进行微调时，可以冻结这部分参数不跟随训练。当你设置require_grad=False时，你的这部分参数将不会在前向传播的过程中不会被保存记录，他们不会在后向传递中更新他们的 .grad字段，因为它们一开始就不会成为后向图的一部分。因为这是一种常见的模式，所以也可以使用 nn.Module.requires_grad_() 在模块级别设置requires_grad。当应用于模块时， .requires_grad_() 会影响模块的所有参数（默认情况下 requires_grad=True ）。

梯度模式

除了设置 requires_grad之外，还可以从 Python 启用三种可能的模式，这些模式可以影响 PyTorch 中 autograd 计算的处理的方式：默认模式（grad 模式）、no-grad 模式和inference模式，所有这些都可以通过context manager和decorator。

默认模式

“默认模式”实际上是我们在没有启用其他模式（如 no-grad 和 inference 模式）时隐含的模式。与“no-grad mode”相比，默认模式有时也称为“grad mode”。
关于默认模式，最重要的一点是它是 requires_grad 生效的唯一模式。在其他两种模式下，requires_grad 总是被覆盖为 False。

no-grad模式

在no-grad模式下，所有的计算都表现为所有的输入都不需要有梯度，即使有 require_grad=True 的输入，no-grad 模式下的计算也永远不会记录在后向图中。
当您需要执行不应该由 autograd 记录的操作时启用 no-grad 模式，但您以后仍希望在 grad 模式下使用这些计算的输出。这个context-manager可以方便地禁用代码块或函数的梯度，而无需临时将张量设置为 requires_grad=False。

例如，在编写优化器时，no-grad 模式可能很有用：在执行训练更新时，您希望就地更新参数，而 autograd 不会记录更新。您还打算在下一次正向传递中使用更新的参数以 grad 模式进行计算。

torch.nn.init 中的实现在初始化参数时也依赖于 no-grad 模式，以避免在就地更新初始化参数时进行 autograd 跟踪。

Inference 模式

推理模式是无梯度模式的极端版本​​。就像在 no-grad 模式下一样，推理模式下的计算不会记录在后向图中，但启用推理模式将允许 PyTorch 进一步加速您的模型。这种更好的运行时有一个缺点：在推理模式下创建的张量将无法用于退出推理模式后由 autograd 记录的计算。

当您执行不需要在后向图中记录的计算时启用推理模式，并且您不打算在稍后将由 autograd 记录的任何计算中使用在推理模式中创建的张量。

建议您在代码中不需要自动梯度跟踪的部分（例如，数据处理和模型评估）尝试推理模式。如果它为您的用例开箱即用，那将是免费的性能胜利。如果您在启用推理模式后遇到错误，请检查您是否没有在退出推理模式后由 autograd 记录的计算中使用在推理模式下创建的张量。如果在您的情况下无法避免此类使用，您可以随时切换回 no-grad 模式。

评估(Evaluation)模式 (nn.Module.eval())

评估模式实际上并不是一种在本地禁用梯度计算的机制。在这里为什么要介绍它呢，因为它有时会被混淆为这样的机制。

在功能上，module.eval()（或等效的 module.train()）与 no-grad 模式和推理模式完全正交(即不存在相关性)。 model.eval() 如何影响您的模型完全取决于模型中使用的特定模块以及它们是否定义任何训练模式特定行为。

如果您的模型依赖于诸如 torch.nn.Dropout 和 torch.nn.BatchNorm2d 之类的模块，这些模块可能会因训练模式而有所不同，则您应当区分调用 model.eval() 和 model.train()。
例如，为了避免更新验证数据的 BatchNorm 运行统计信息。

建议您在训练时始终使用 model.train()，在评估模型（验证/测试）时始终使用 model.eval()，即使您不确定模型是否具有特定于训练模式的行为，因为您使用的某个模块可能会更新为在训练和评估模式下表现不同。

使用 autograd 进行原地操作

对于使用in-place的操作运算，pytorch这边不是非常建议使用，Autograd本身的缓冲区的释放和重用的效率是非常高的，极少情况下进行原地操作会显著降低内存的使用，

限制原地操作有两个主要原因：

  1. 原地操作可能会覆盖计算梯度所需要的值。
  2. 原地操作需要重写计算图的操作实现，out-of-place的方法只需要分配新的内存给创造的对象，并且保存了对旧图的引用。在原地操作时，需要将所有输入的创建者更改为表示此操作的函数。这可能很棘手，特别是如果有许多张量引用相同的存储。比如，通过索引或转置创建。如果修改的输入的存储被任何其他张量引用，则in-place函数实际上会引发错误。

pytorch如何保证原地操作的正确性？

每个张量都有一个版本计数器，每次它在任何操作中被标记为脏时都会递增。当一个函数保存任何用于后向计算的张量时，它们包含的张量的版本计数器也会被保存。当你访问self.saved_tensor时它就会被检查，如果它大于保存的值，则会引发错误。这可以确保如果您使用原地函数并且没有看到任何错误，则可以确保计算出的梯度是正确的。

多线程 Autograd

autograd 引擎负责运行计算反向传递所需的所有反向操作。下面将介绍在多线程环境中充分利用它的所有细节。
（仅与 PyTorch 1.6+ 相关，因为以前版本中的行为不同）

用户可以使用多线程代码（例如 Hogwild 训练）训练他们的模型，并且不会阻塞并发的反向计算，示例代码可以是：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

# Define a train function to be used in different threads
def train_fn():
    x = torch.ones(5, 5, requires_grad=True)
    # forward
    y = (x + 3) * (x + 4) * 0.5
    # backward
    y.sum().backward()
    # potential optimizer update


# User write their own threading code to drive the train_fn
threads = []
for _ in range(10):
    p = threading.Thread(target=train_fn, args=())
    p.start()
    threads.append(p)

for p in threads:
    p.join()

一些需要被了解的特性：

CPU 上的并发

当您在 CPU 上的多个线程中通过 python 或 C++ API 运行backward() 或 grad() 时，将会看到额外的并发性，而不是在执行期间以特定顺序序列化所有向后调用（PyTorch 1.6 之前的行为）。

Non-determinism

如果您同时在多个线程上调用backward()，但使用共享输入（即Hogwild CPU 训练）。由于参数在线程之间自动共享，梯度累积可能在线程之间的反向调用中变得不确定，因为两个反向调用可能会访问并尝试累积相同的 .grad属性。这在技术上是不安全的，它可能会导致赛车状态，结果可能无法使用。

但是，如果您使用多线程方法来驱动整个训练过程但使用共享参数，那么这是预期的模式，使用多线程的用户应该牢记线程模型并且应该预期会发生这种情况。用户可以使用功能 API torch.autograd.grad() 来计算梯度，而不是使用 backward() 来避免不确定性。

Graph retaining

如果autograd图结构在多线程中被共享，例如图的第一部分运行在单线程上，图的第二部分采用多线程进行优化，在这种情况下，图结构的第一部分被优化，在这种情况下，不同的线程在同一个图上执行 grad()或 backward()可能会在一个线程的运行中破坏图，并且在这种情况下另一个线程将崩溃。 Autograd 将向用户发出错误，类似于在没有 retain_graph=True 的情况下两次调用 backward()，并让用户知道他们应该使用 retain_graph=True。

在 Autograd节点上的线程安全。

由于 Autograd 允许调用者线程驱动其向后执行以实现潜在的并行性，因此我们必须确保 CPU 上的线程安全，并行反向传播共享 GraphTask 的部分/全部。

由于 GIL，自定义 Python autograd.function 是自动线程安全的。对于内置 C++ Autograd 节点（例如 AccumulateGrad、CopySlices）和自定义 autograd::Function，Autograd 引擎使用线程互斥锁来保护可能具有状态写入/读取的 autograd 节点上的线程安全。

C++ hooks上没有线程安全

Autograd 依赖于用户编写线程安全的 C++ hooks。如果您希望hooks在多线程环境中正确应用，则需要编写适当的线程锁定代码以确保hooks是线程安全的。

对于复数的Autograd

简要概括:

  1. 当你对具有复数域或者共域的函数$f(z)$进行微分时，梯度是在假设函数是更大的实值损失函数$g(input)=L$的一部分的情况下计算的，
    $\frac{\partial L}{\partial z^{}}$($z^{}$代表 $z$ 的共轭)。其负值正是梯度下降算法中使用的最陡下降方向。因此说有的优化器对于复数参数都是开箱即用的。
  2. 此约定与 TensorFlow 的复杂微分约定相匹配，但与 JAX(计算$\frac{\partial L}{\partial z}$) 不同。
  3. 如果您有一个内部使用复数操作的实对实函数，则此处的约定无关紧要：您将始终获得与仅使用实数操作实现相同的结果。
O(∩_∩)O哈哈~