Autograd 机制 ¶
译者:片刻小哥哥
本笔记将概述 autograd 的工作原理并记录操作。并不是绝对有必要理解所有这些,但我们建议您熟悉它,因为它将帮助您编写更高效、更简洁的程序,并可以帮助您进行调试。
autograd 如何编码历史记录 ¶
Autograd 是一种反向自动微分系统。从概念上讲,autograd 记录了一个图,记录了执行操作时创建数据的所有操作,为您提供一个有向无环图,其叶子是输入tensor,根是输出tensor。通过从根到叶子追踪该图,您可以自动计算使用链式法则的梯度。
在内部,autograd 将此图表示为“Function”对象(实际上是表达式)的图,可以对其进行“apply()”编辑以计算评估该图的结果。在计算前向传递时,autograd 同时执行请求的计算并构建一个表示计算梯度的函数的图(每个 torch.Tensor
是该图的入口点)。当前向传递完成时,我们在后向传递中评估该图以计算梯度。
需要注意的重要一点是,图形在每次迭代时都会从头开始重新创建,这正是允许使用任意 Python 控制流语句的原因,这些语句可以在每次迭代时更改图形的整体形状和大小。在启动训练之前,您不必对所有可能的路径进行编码 - 您运行的内容就是您所区分的。
保存的tensor ¶
某些操作需要在前向传递期间保存中间结果,以便执行后向传递。例如,函数 \(x ↦ x^2\) 保存输入 \(x\) 来计算梯度。
定义自定义 Python Function
时,可以使用 save_for_backward()
在前向传递过程中保存tensor和saved_tensors
以在向后传递期间检索它们。有关更多信息,请参阅扩展 PyTorch。
对于 PyTorch 定义的操作(例如 torch.pow()
),tensor会根据需要自动保存。您可以通过查找以前缀“_saved”开头的属性来探索(出于教育或调试目的)某个“grad_fn”保存了哪些tensor。
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 对象。但情况可能并不总是如此。例如:
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“打包”了保存的tensor,并将其“解包”到不同的tensor中以供读取。在这里,您通过访问“y.grad_fn._saved_result”获得的tensor是与“y”不同的tensor对象(但它们仍然共享相同的存储)。
tensor是否会被打包到不同的tensor对象中取决于它是否是其自己的 grad_fn 的输出,这是一个可能会更改的实现细节,用户不应依赖。
您可以使用 保存tensor的钩子 控制 PyTorch 如何打包/解包。
不可微函数的梯度 ¶
使用自动微分的梯度计算仅在使用的每个初等函数可微时才有效。不幸的是,我们在实践中使用的许多函数不具有此属性(例如,“0”处的“relu”或“sqrt”)。为了尝试减少不可微函数的影响,我们通过按顺序应用以下规则来定义基本运算的梯度:
- 如果函数可微,因此当前点存在梯度,请使用它。
- 如果函数是凸函数(至少是局部凸函数),则使用最小范数的次梯度(即最速下降方向)。
- 如果函数是凹函数(至少是局部凹函数),则使用最小范数的超梯度(考虑 -f(x) 并应用前一点)。
- 如果定义了函数,则通过连续性定义当前点的梯度(请注意,此处可以使用“inf”,例如“sqrt(0)”)。如果可能有多个值,请任意选择一个。
- 如果函数未定义(例如
sqrt(-1)
、log(-1)
或输入为NaN
时的大多数函数),则用作梯度的值是任意的(我们也可以提出一个错误,但不能保证)。大多数函数将使用“NaN”作为梯度,但出于性能原因,某些函数将使用其他值(例如“log(-1)”)。 - 如果函数不是确定性映射(即它不是数学函数),它将被标记为不可微分。如果在“no_grad”环境之外需要 grad 的tensor上使用,这将使其向后出错。
本地禁用梯度计算 ¶
Python 有多种机制可用于本地禁用梯度计算:
要禁用整个代码块的梯度,可以使用上下文管理器,例如无梯度模式和推理模式。为了更细粒度地从梯度计算中排除子图,可以设置tensor的“requires_grad”字段。
下面,除了讨论上面的机制之外,我们还描述了评估模式(nn.Module.eval()
),这种方法不用于禁用梯度计算,但由于其名称,经常与这三种方法混淆。
设置 requires_grad
¶
requires_grad
是一个标志,默认为 false *除非包装在
* nn.Parameter
中,它允许从梯度计算中细粒度地排除子图。它在前向传播和后向传播中都有效:
在前向传递过程中,如果至少有一个输入tensor需要 grad,则该操作仅记录在后向图中。在后向传递过程中( .backward()
),只有具有 requires_grad=True
的叶tensor才会有梯度累积到它们的“.grad”字段中。
值得注意的是,即使每个tensor都有这个标志,设置它只对叶tensor有意义(没有 grad_fn
的tensor,例如 nn.Module
的参数)。非叶tensor(具有 grad_fn
的tensor)是具有与之关联的后向图的tensor。因此,它们的梯度将需要作为中间结果来计算需要梯度的叶tensor的梯度。从这个定义可以清楚地看出,所有非叶tensor将自动具有“require_grad=True”。
设置“requires_grad”应该是控制模型的哪些部分参与梯度计算的主要方式,例如,如果您需要在模型微调期间冻结预训练模型的部分。
要冻结模型的某些部分,只需将“.requires_grad_(False)”应用于您不想更新的参数即可。如上所述,由于使用这些参数作为输入的计算不会记录在前向传递中,因此它们的“.grad”字段不会在后向传递中更新,因为它们不会成为第一个中后向图的一部分地点,随心所欲。
因为这是一种常见的模式,所以也可以使用 nn.Module.requires_grad_()
在模块级别设置 requires_grad
。当应用于模块时,.requires_grad_()
对模块的所有参数生效(默认情况下具有 requires_grad=True
)。
毕业模式 ¶
除了设置 requires_grad
之外,还可以从 Python 中选择三种 grad 模式,这些模式会影响 PyTorch 中的计算在内部如何由 autograd 处理:默认模式(grad 模式)、no-grad 模式和推理模式,所有这些模式可以通过上下文管理器和装饰器进行切换。
Mode | Excludes operations from being recorded in backward graph | Skips additional autograd tracking overhead | Tensors created while the mode is enabled can be used in grad-mode later | Examples |
---|---|---|---|---|
default | ✓ | Forward pass | ||
no-grad | ✓ | ✓ | Optimizer updates | |
inference | ✓ | ✓ | Data processing, model evaluation |
默认模式(渐变模式)¶
“默认模式”是当没有启用其他模式(例如无梯度和推理模式)时我们隐式处于的模式。为了与“无渐变模式”对比,默认模式有时也称为“渐变模式”。
关于默认模式,最重要的是要了解它是 requires_grad
生效的唯一模式。在其他两种模式中,“requires_grad”始终被覆盖为“False”。
无梯度模式 ¶
无梯度模式下的计算行为就像没有输入需要梯度一样。换句话说,无梯度模式下的计算永远不会记录在后向图中,即使存在具有 require_grad=True
的输入。
当您需要执行不应该由自动分级记录的操作时,请启用无分级模式,但您仍希望稍后在分级模式下使用这些计算的输出。这个上下文管理器可以方便地禁用代码块或函数的梯度,而无需临时将tensor设置为“requires_grad=False”,然后返回“True”。
例如,在编写优化器时,无梯度模式可能很有用:当执行训练更新时,您希望就地更新参数,而不用自动梯度记录更新。您还打算在下一个向前传球。
torch.nn.init 中的实现在初始化参数时也依赖于无梯度模式,以避免在就地更新初始化参数时进行自动梯度跟踪。
推理模式 ¶
推理模式是无梯度模式的极端版本。就像在无梯度模式中一样,推理模式下的计算不会记录在后向图中,但是启用推理模式将使 PyTorch 能够进一步加速您的模型。这种更好的运行时有一个缺点:在推理模式下创建的tensor将无法用于退出推理模式后由 autograd 记录的计算。
当您执行不需要记录在后向图中的计算时,启用推理模式,并且您不打算在稍后由 autograd 记录的任何计算中使用推理模式中创建的tensor。
建议您在代码中不需要自动分级跟踪的部分尝试推理模式(例如,数据处理和模型评估)。如果它对于您的用例来说是开箱即用的,那么它就是免费的性能胜利。如果您在启用推理模式后遇到错误,请检查您是否未在退出推理模式后由 autograd 记录的计算中使用在推理模式中创建的tensor。如果您无法避免在您的情况下使用这种情况,您可以随时切换回无梯度模式。
有关推理模式的详细信息请参见推理模式。
有关推理模式的实现细节,请参阅 RFC-0011-InferenceMode 。
评估模式 ( nn.Module.eval()
) ¶
评估模式不是一种局部禁用梯度计算的机制。无论如何,它都包含在这里,因为它有时会被混淆为这样一种机制。
从功能上讲, module.eval() (或等效的 module.train(False) )与无梯度模式和推理模式完全正交。 model.eval() 如何影响您的模型完全取决于模型中使用的特定模块以及它们是否定义任何训练模式特定行为。
如果您的模型依赖于 torch.nn.Dropout
和torch.nn.BatchNorm2d
根据训练模式,其行为可能会有所不同,例如,避免更新验证数据上的 BatchNorm 运行统计数据。
建议您在训练时始终使用 model.train() ,在评估模型(验证/测试)时始终使用 model.eval() ,即使您不确定模型是否具有训练模式特定行为,因为您的模块正在使用的可能会更新为在训练和评估模式下表现不同。
使用 autograd 进行就地操作 ¶
支持 autograd 中的就地操作是一件困难的事情,我们在大多数情况下不鼓励使用它们。 Autograd 积极的缓冲区释放和重用使其非常高效,并且很少有就地操作显着降低内存使用量的情况。除非您在内存压力很大的情况下运行,否则您可能永远不需要使用它们。
限制就地操作的适用性的主要原因有两个:
- 就地操作可能会覆盖计算梯度所需的值。2.每个就地操作都需要实现重写计算图。异地版本只是分配新对象并保留对旧图的引用,而就地操作则需要将所有输入的创建者更改为表示此操作的“函数”。这可能很棘手,特别是如果有许多tensor引用相同的存储(例如通过索引或转置创建),并且如果修改后的输入的存储被任何其他“tensor”引用,则就地函数将引发错误。
就地正确性检查 ¶
每个tensor都保留一个版本计数器,每次在任何操作中被标记为脏时,该计数器都会递增。当函数保存任何向后tensor时,也会保存其包含tensor的版本计数器。一旦您访问 self.saved_tensors ,就会对其进行检查,如果它大于保存的值,则会引发错误。这确保了如果您使用就地函数并且没有看到任何错误,则可以确定计算的梯度是正确的。
多线程 Autograd ¶
autograd 引擎负责运行计算反向传递所需的所有反向操作。本节将描述可以帮助您在多线程环境中充分利用它的所有细节。 (这仅与 PyTorch 1.6 + 相关,因为之前版本的行为有所不同。)
用户可以使用多线程代码(例如 Hogwild 训练)来训练他们的模型,并且不会阻止并发向后计算,示例代码可以是:
# 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 之前的行为)。
非决定论 ¶
如果您同时从多个线程调用“backward()”并共享输入(即 Hogwild CPU 训练),那么应该预期不确定性。这种情况可能会发生,因为参数会在线程之间自动共享,因此,多个线程可能会访问和尝试在梯度累积期间累积相同的“.grad”属性。这在技术上并不安全,并且可能会导致竞争条件,并且结果可能无法使用。
开发具有共享参数的多线程模型的用户应该牢记线程模型,并且应该理解上述问题。
函数式 API torch.autograd.grad()
可用于计算梯度,而不是 back()
以避免非确定性。
图形保留 ¶
如果自动分级图的一部分在线程之间共享,即运行前向单线程的第一部分,然后在多个线程中运行第二部分,则图的第一部分被共享。在这种情况下,不同的线程在同一个图上执行“grad()”或“backward()”可能会出现一个线程动态销毁图的问题,而在这种情况下另一个线程将崩溃。 Autograd 会向用户发出类似于两次调用“backward()”而没有使用“retain_graph=True”的错误,并让用户知道他们应该使用“retain_graph=True”。
Autograd 节点上的线程安全 ¶
由于 Autograd 允许调用者线程驱动其向后执行以实现潜在的并行性,因此我们必须通过共享部分/全部 GraphTask 的并行“backward()”调用来确保 CPU 上的线程安全。
由于 GIL,自定义 Python autograd.Function
自动是线程安全的。对于内置 C++ Autograd 节点(例如 AccumulateGrad、CopySlices)和自定义 autograd::Function
,Autograd 引擎使用线程互斥锁来确保线程安全autograd 可能具有写入/读取状态的节点。
C++ 钩子上没有线程安全 ¶
Autograd 依赖于用户编写线程安全的 C++ 钩子。如果希望钩子能够正确应用在多线程环境中,则需要编写适当的线程锁定代码以确保钩子是线程安全的。
复数的 Autograd ¶
简短版本:
- 当你使用PyTorch来微分任何函数时 \(f(z)\) 对于复杂域和/或共域,梯度是在函数是较大实值损失函数的一部分的假设下计算的 \(g(input)=L\)。计算出的梯度为 \frac{∂L}{∂z*} (注意 z 的共轭),其负数正是梯度下降算法中使用的最速下降方向。因此,所有现有的优化器都可以使用复杂的参数进行开箱即用。
- 此约定与 TensorFlow 的复杂微分约定相匹配,但与 JAX 不同(JAX 计算 \(\frac{∂L}{∂z}\) ).
- 如果您有一个内部使用复杂操作的实数到实数函数,那么这里的约定并不重要:您将始终获得与仅使用实数操作实现的结果相同的结果。
如果您对数学细节感到好奇,或者想知道如何在 PyTorch 中定义复杂的导数,请继续阅读。
什么是复杂导数? ¶
复数可微分的数学定义采用导数的极限定义,并将其推广到对复数进行运算。考虑一个函数 \(f: ℂ → ℂ\),
在哪里 \(u\) 你和 \(v\) 是两个变量实值函数并且 \(j\) 是虚数单位。
使用导数定义,我们可以写:
为了使这个极限存在,不仅必须 \(u\) 你和 \(v\) 必须是实数可微,但是 \(f\) 还必须满足柯西-黎曼方程。 换句话说:为实部和虚部计算的极限 \((f)\) 必须相等。 这是一个更严格的条件。
复可微函数通常称为全纯函数。它们表现良好,具有您从真正的可微函数中看到的所有良好属性,但在优化世界中实际上毫无用处。对于优化问题,研究界仅使用实值目标函数,因为复数不属于任何有序域,因此具有复值损失没有多大意义。
事实证明,没有有趣的实值目标满足柯西-黎曼方程。所以同态函数的理论不能用于优化,因此大多数人使用Wirtinger演算。
维廷格微积分出现了…… ¶
所以,我们有复可微分和全纯函数的伟大理论,但我们根本无法使用其中的任何一个,因为许多常用的函数都不是全纯的。一个贫穷的数学家该怎么办?好吧,维廷格观察到,即使
f(z)
不是全纯的,可以将其重写为二变量函数
f(z,z*)
这总是全纯的。这是因为实部和虚部的组成部分
z
可以表示为
z
和
z^*
作为:
R
e
(
z
)
=
z
+
z
*
2
I
m
(
z
)
=
z
−
z
*
2
j
egin{对齐} \mathrm{Re}(z) &= rac {z + z^}{2} \ \mathrm{Im}(z) &= rac {z - z^}{2j} \end{对齐}
重新
(
z
)
我
(
z
)
2
+
与
*
=
2
j
z
-
*
维廷格微积分建议学习
f ( z ,
z*
)
f(z, z^*)
f ( z ,
z
*
)
相反,如果
F
F
F
是真正可微分的(另一种方式认为它是坐标系的变化,从
f ( x , y )
f(x,y)
f ( x ,
)
到
f ( z ,
z*
)
f(z, z^*)
f ( z ,
z
*
)
.) 该函数具有偏导数
∂
∂z
rac{\partial }{\partial z}
∂z
∂
和
∂
∂
z*
rac{\partial}{\partial z^{*}}
∂
z
*
∂
我们可以使用链式法则来建立这些偏导数和偏导数之间的关系,即
z
z
z
。
∂
∂
x
=
∂
z
∂
x
∗
∂
∂
z
+
∂
z
∗
∂
x
∗
∂
∂
z
∗
=
∂
∂
z
+
∂
∂
z
∗
∂
∂
y
=
∂
z
∂
y
∗
∂
∂
z
+
∂
z
∗
∂
y
∗
∂
∂
z
∗
=
1
j
∗
(
∂
∂
z
−
∂
∂
z
∗
)
egin{aligned} rac{\partial }{\partial x} &= rac{\partial z}{\partial x} * rac{\partial }{\partial z} + rac{\partial z^}{\partial x} * rac{\partial }{\partial z^} \ &= rac{\partial }{\partial z} + rac{\partial }{\partial z^} \ \ rac{\partial }{\partial y} &= rac{\partial z}{\partial y} * rac{\partial }{\partial z} + rac{\partial z^}{\partial y} * rac{\partial }{\partial z^} \ &= 1j * \left(rac{\partial }{\partial z} - rac{\partial }{\partial z^} ight) \end{aligned}
∂
x
∂
∂
y
∂
=
∂
x
∂
z
∗
∂
z
∂
+
∂
x
∂
z
∗
∗
∂
z
∗
∂
=
∂
z
∂
+
∂
z
∗
∂
=
∂
y
∂
z
∗
∂
z
∂
+
∂
y
∂
z
∗
∗
∂
z
∗
∂
=
1
j
∗
(
∂
z
∂
−
∂
z
∗
∂
)
从上面的方程,我们得到:
∂
∂
z
=
1
/
2
*
(
∂
∂
x
−
1
j
*
∂
∂
y
)
∂
∂
z
*
=
1
/
2
*
(
∂
∂
x
+
1
j
∗
∂
∂
y
)
egin{aligned} rac{\partial }{\partial z } &= 1/2 * \left(rac{\partial }{\partial x} - 1j * rac{\partial }{\partial y} ight) \ rac{\partial }{ \partial z^*} &= 1/2 * \left(rac{\partial }{\partial x} + 1j * rac{\partial }{\partial y} ight) \end{对齐}
∂
z
∂
∂
z
*
∂
=
1/2
*
(
∂
x
∂
−
1
j
*
∂
y
∂
)
=
1/2
*
(
∂
x
∂
+
1
j
*
∂
y
∂
)
这是维廷格微积分的经典定义,您可以在维基百科上找到。
这一变化带来了很多美好的后果。
- 其一,柯西-黎曼方程可以简单地解释为:
∂f
∂
z*
= 0
rac{\partial f}{\partial z^*} = 0
∂
z
*
∂f
=
0
(也就是说,函数
F
F
F
可以完全写成
z
z
z
,没有参考
z*
z^*
z
*
). * 另一个重要的(有点违反直觉的)结果,正如我们稍后将看到的,是当我们对实值损失进行优化时,我们在进行变量更新时应该采取的步骤由下式给出
∂ 损失
∂
z*
rac{\部分损失}{\部分z^*}
∂
z
*
∂ 损失
(不是
∂ 损失
∂z
rac{\部分损失}{\部分z}
∂z
∂ 损失
)。
如需更多阅读,请查看:https://arxiv.org/pdf/0906.4835.pdf
维廷格微积分在优化中有何用处? ¶
音频和其他领域的研究人员更常见地使用梯度下降来优化具有复杂变量的实值损失函数。通常,这些人将实值和虚值视为可以更新的单独通道。对于步长
A2
\α/2
A2
和损失
L
L
L
,我们可以写出以下方程
R 2
ℝ^2
右
2
:
x
n
+
1
=
x
n
−
(
α
/
2
)
∗
∂
L
∂
x
y
n
+
1
=
y
n
−
(
α
/
2
)
*
∂
L
∂
y
egin{对齐} x_{n+1} &= x_n - (lpha/2) * rac{\partial L}{\partial x} \ y_{n+1} &= y_n - (lpha/2) * rac{\partial L}{\partial y} \end{对齐}
x
n
+
1
y
n
+
1
=
x
n
−
(
α
/2
)
∂
∂
x
∂
L
=
y
n
-
(
α
/2
)
*
∂
y
∂
L
这些方程如何转化为复杂空间
C
ℂ
C
?
z
n
+
1
=
x
n
−
(
a
/
2
)
∗
∂
L
∂
x
+
1
j
*
(
y
n
−
(
a
/
2
)
*
∂
L
∂
y
)
=
z
n
−
a
*
1
/
2
*
(
∂
L
∂
x
+
j
∂
L
∂
y
)
=
z
n
−
a
*
∂
L
∂
z
∗
egin{aligned} z_{n+1} &= x _n - (lpha/2) * rac{\partial L}{\partial x} + 1j * (y_n - (lpha/2) * rac{\partial L}{\partial y}) \ &= z_n - lpha * 1/2 * \left(rac{\partial L}{\partial x} + j rac{\partial L}{\partial y } ight) \ &= z_n - lpha * rac{\partial L}{\partial z^*} \end{对齐}
z
n
+
1
x
n
−
(
a
/2
)
*
∂
x
∂
L
+
1
j
*
(
y
n
−
(
a
/2
)
*
∂
y
∂
L
)
=
z
n
-
a
*
1/2
*
(
∂
x
∂
L
+
j
∂
y
∂
L
)
=
z
n
-
a
*
∂
z
*
∂
L
非常有趣的事情发生了:维廷格微积分告诉我们,我们可以将上面的复变量更新公式简化为仅指共轭维廷格导数
∂L
∂
z*
rac{\部分L}{\部分z^*}
∂
z
*
∂L
,给我们准确的优化步骤。
由于共轭 Wirtinger 导数为我们提供了实值损失函数的正确步骤,因此当您对函数与实值损失进行微分时,PyTorch 会为您提供该导数。
PyTorch 如何计算共轭 Wirtinger 导数? ¶
通常,我们的导数公式将 grad_output 作为输入,表示我们已经计算出的传入向量雅可比积,又名,
∂L
∂
s*
rac{\partial L}{\partial s^*}
∂
s
*
∂L
, 在哪里
L
L
L
是整个计算的损失(产生实际损失)并且
s
s
s
是我们函数的输出。这里的目标是计算
∂L
∂
z*
rac{\部分L}{\部分z^*}
∂
z
*
∂L
, 在哪里
z
z
z
是函数的输入。事实证明,在实际损失的情况下,我们可以仅计算
∂L
∂
s*
rac{\partial L}{\partial s^*}
∂
s
*
∂L
,尽管链式法则意味着我们还需要访问
∂L
∂s
rac{\partial L}{\partial s}
∂s
∂L
。如果您想跳过此推导,请查看本节中的最后一个方程,然后跳到下一节。
让我们继续合作
f : C → C
f: ℂ → ℂ
F
:
C
→
C
定义为
f ( z ) = f ( x + y j ) = u ( x , y ) + v ( x , y ) j
f(z) = f(x+yj) = u(x, y) + v(x, y)j
f(z)
=
f ( x
+
yj)
=
你(x,
)
+
v ( x ,
) j
。如上所述,autograd 的梯度约定以实值损失函数的优化为中心,所以我们假设
F
F
F
是较大实值损失函数的一部分
G
G
G
。使用链式法则,我们可以写出:
(1) ¶
∂
L
∂
z
*
=
∂
L
∂
u
*
∂
u
∂
z
*
+
∂
L
∂
v
*
∂
v
∂
z
∗
rac{\partial L}{\partial z^} = rac{\partial L }{\partial u} * rac{\partial u}{\partial z^} + rac{\partial L}{\partial v} * rac{\partial v}{\partial z^ *}
∂
z
*
∂
L
=
∂
u
∂
L
*
∂
z
*
∂
u
+
∂
v
∂
L
*
∂
z
*
∂
v
现在使用 Wirtinger 导数定义,我们可以写:
∂
L
∂
s
=
1
/
2
*
(
∂
L
∂
u
−
∂
L
∂
v
j
)
∂
L
∂
s
*
=
1
/
2
*
(
∂
L
∂
u
+
∂
L
∂
v
j
)
egin{aligned} rac {\partial L}{\partial s} = 1/2 * \left(rac{\partial L}{\partial u} - rac{\partial L}{\partial v} j ight) \ rac{\partial L}{\partial s^*} = 1/2 * \left(rac{\partial L}{\partial u} + rac{\partial L}{\partial v} j ight) \end{对齐}
∂
s
∂
L
=
1/2
*
(
∂
u
∂
L
-
∂
v
∂
L
j
)
∂
s
*
∂
L
=
1/2
*
(
∂
u
∂
L
+
∂
v
∂
L
j
)
这里需要注意的是,由于
你
你
你
和
v
v
v
是实函数,并且
L
L
L
根据我们的假设是真实的
F
F
F
是实值函数的一部分,我们有:
(2) ¶
(
∂
L
∂
s
)
*
=
∂
L
∂
s
*
\left( rac{\partial L}{\partial s} ight)^ = rac{\partial L}{\partial s^}
(
∂
s
∂
L
)
*
=
∂
s
*
∂
L
IE。,
∂L
∂s
rac{\partial L}{\partial s}
∂s
∂L
等于
梯度 _ 输出 p u
t*
梯度_输出^*
毕业_ o u tp u
t
*
。
求解上述方程
∂L
∂你
rac{\偏L}{\偏u}
∂你
∂L
和
∂L
∂v
rac{\部分L}{\部分v}
∂v
∂L
,我们得到:
(3) ¶
∂
L
∂
u
=
∂
L
∂
s
+
∂
L
∂
s
∗
∂
L
∂
v
=
−
1
j
*
(
∂
L
∂
s
−
∂
L
∂
s
∗
)
egin{对齐} rac{\partial L}{ \partial u} = rac{\partial L}{\partial s} + rac{\partial L}{\partial s^} \ rac{\partial L}{\partial v} = - 1j * \left(rac{\partial L}{\partial s} - rac{\partial L}{\partial s^} ight) \end{对齐}
∂
u
∂
L
=
∂
s
∂
L
+
∂
s
*
∂
L
∂
v
∂
L
=
−
1
j
∗
(
∂
s
∂
L
−
∂
s
*
∂
L
)
∂
L
∂
z
∗
=
(
∂
L
∂
s
+
∂
L
∂
s
∗
)
∗
∂
u
∂
z
∗
−
1
j
∗
(
∂
L
∂
s
−
∂
L
∂
s
∗
)
∗
∂
v
∂
z
∗
=
∂
L
∂
s
∗
(
∂
u
∂
z
∗
+
∂
v
∂
z
∗
j
)
+
∂
L
∂
s
∗
∗
(
∂
u
∂
z
∗
−
∂
v
∂
z
∗
j
)
=
∂
L
∂
s
∗
∗
∂
(
u
+
v
j
)
∂
z
∗
+
∂
L
∂
s
∗
∂
(
u
+
v
j
)
∗
∂
z
∗
=
∂
L
∂
s
∗
∂
s
∂
z
∗
+
∂
L
∂
s
∗
∗
∂
s
∗
∂
z
∗
egin{aligned} rac{\partial L}{\partial z^} &= \left(rac{\partial L}{\partial s} + rac{\partial L}{\partial s^} ight) * rac{\partial u}{\partial z^} - 1j * \left(rac{\partial L}{\partial s} - rac{\partial L}{\partial s^} ight) * rac{\partial v}{\partial z^} \ &= rac{\partial L}{\partial s} * \left(rac{\partial u}{\partial z^} + rac{\partial v}{\partial z^} j ight) + rac{\partial L}{\partial s^} * \left(rac{\partial u}{\partial z^} - rac{\partial v}{\partial z^} j ight) \ &= rac{\partial L}{\partial s^} * rac{\partial (u + vj)}{\partial z^} + rac{\partial L}{\partial s} * rac{\partial (u + vj)^}{\partial z^} \ &= rac{\partial L}{\partial s} * rac{\partial s}{\partial z^} + rac{\partial L}{\partial s^} * rac{\partial s^}{\partial z^} \ \end{aligned}
∂
z
∗
∂
L
=
(
∂
s
∂
L
+
∂
s
∗
∂
L
)
∗
∂
z
∗
∂
u
−
1
j
∗
(
∂
s
∂
L
−
∂
s
∗
∂
L
)
∗
∂
z
∗
∂
v
=
∂
s
∂
L
∗
(
∂
z
∗
∂
u
+
∂
z
∗
∂
v
j
)
+
∂
s
∗
∂
L
∗
(
∂
z
∗
∂
u
−
∂
z
∗
∂
v
j
)
=
∂
s
∗
∂
L
∗
∂
z
∗
∂
(
u
+
v
j
)
+
∂
s
∂
L
∗
∂
z
∗
∂
(
u
+
v
j
)
∗
=
∂
s
∂
L
∗
∂
z
∗
∂
s
+
∂
s
∗
∂
L
∗
∂
z
∗
∂
s
∗
使用 (2) ,我们得到:
(4) ¶
∂
L
∂
z
*
=
(
∂
L
∂
s
*
)
*
*
∂
s
∂
z
*
+
∂
L
∂
s
*
*
(
∂
s
∂
z
)
*
=
(
g
r
a
d
_
o
u
t
p
u
t
)
*
*
∂
s
∂
z
*
+
g
r
a
d
_
o
u
t
p
u
t
*
(
∂
s
∂
z
)
∗
egin{aligned} rac {\partial L}{\partial z^} &= \left(rac{\partial L}{\partial s^} ight)^ * rac{\partial s}{\部分 z^} + rac{\partial L}{\partial s^} * \left(rac{\partial s}{\partial z} ight)^ \ &= oxed{ (grad_output)^ * rac{\partial s}{\partial z^} + grad_output * \left(rac{\partial s}{\partial z } ight)^* } \ \end{对齐}
∂>>
z
*
∂
L
=
(
∂
s
*
∂
L
)
*
*
∂
z
*
∂
s
+
∂
s
*
∂
L
*
(
∂
z
∂
s
)
*
=
(
g
r
a
d
_
o
u
tp>> u
t
)
*
∂
z
*
∂
s
+
g
r
a
d
_
o
u
tp
u
t
*
(
∂
z
∂
s
)
*
最后一个方程对于编写自己的梯度非常重要,因为它将我们的导数公式分解为一个易于手动计算的更简单的公式。
如何为复杂函数编写自己的导数公式? ¶
上面的盒式方程给出了复函数所有导数的一般公式。然而,我们仍然需要计算
∂s
∂z
rac{\部分s}{\部分z}
∂z
∂s
和
∂s
∂
z*
rac{\partial s}{\partial z^*}
∂
z
*
∂s
.有两种方法可以做到这一点:
- 第一种方法是直接使用维廷格导数的定义来计算
∂
s
∂
z
rac{\partial s}{\partial z}
∂
z
∂
s
∂
s
∂
z
*
rac{\partial s}{\partial z^*}
∂
z
*
∂
s
使用
∂
s
∂
x
rac{\partial s}{\partial x}
∂
x
∂
s
和
∂
s
∂
y
rac{\partial s}{\partial y}
∂
y
∂
s
(可以用正常方式计算)。 *第二种方法是使用变量变化技巧和重写
f
(
z
)
f(z)
f
(
z
)
作为二变量函数
f
(
z
,
z
*
)
f(z, z^*)
f
(
z
,
z
∗
)
,并通过处理
z
z
计算 共轭 Wirtinger 导数
z
和
z
*
z^*
z
∗
作为自变量。这通常更容易;例如,如果所讨论的函数是全纯的,则仅使用
z
z
z
(并且>>
∂
s
∂
z
∗
rac{\partial s}{\partial z^*}
∂
z
*
∂
s
将为零)。
我们来考虑一下这个函数
f ( z = x + y j ) = c * z = c * ( x + y j )
f(z = x + yj) = c * z = c * (x+yj)
f ( z
=
X
+
yj)
=
C
∗
z
=
C
*
( X
+
yj)
举个例子,其中
c ∈ R
c \in ℝ
C
ε
右
。
使用第一种方法计算 Wirtinger 导数,我们得到了。
∂s
∂z
= 1 /2 *
(
∂s
∂x
-
∂s
∂y
j)
= 1 /2 * ( c − ( c * 1 j ) * 1 j )
=c
∂s
∂
z*
= 1 /2 *
(
∂s
∂x
+
∂s
∂y
j)
= 1/2(c+(c1j)*1j)
= 0
egin{对齐} rac{\partial s}{\partial z} &= 1/2 * \left(rac{\partial s}{\partial x} - rac{\partial s}{\partial y} j ight) \ &= 1/2 * (c - (c * 1j) * 1j) \ &= c \ \ \ rac{\partial s}{\partial z ^*} &= 1/2 * \left(rac{\partial s}{\partial x} + rac{\partial s}{\partial y} j ight) \ &= 1/2 * (c + (c * 1j) * 1j) \ &= 0 \end{对齐}
∂z
∂s
∂
z
*
∂s
=
1/2
*
(
∂x
∂s
-
∂y
∂s
j
)
=
1/2
*
( C
-
( C
*
1 j )
*
1 j )
=
C
=
1/2
*
(
∂x
∂s
+
∂y
∂s
j
)
=
1/2
*
( C
+
( C
*
1 j )
*
1 j )
=
0
使用 (4) 和 grad_output = 1.0 (这是在 PyTorch 中对标量输出调用 backward()
时使用的默认 grad 输出值),我们得到:
∂
L
∂
z
*
=
1
*
0
+
1
*
c
=
c
\frac{\partial L}{\partial z^*} = 1 * 0 + 1 * c = c
∂
z
*
∂
L
=
1
*
0
+
1
*
c
=
c
使用第二种方法计算 Wirtinger 导数,我们直接得到:
∂
s
∂
z
=
∂
(
c
*
z
)
∂
z
=
c
∂
s
∂
z
∗
=
∂
(
c
*
z
)
∂
z
*
=
0
egin{对齐} rac{\partial s}{\partial z} &= rac {\partial (cz)}{\partial z} \ &= c \ rac{\partial s}{\partial z^} &= rac{\partial (c z)}{\partial z^} \ &= 0 \end{对齐}
∂
z
∂
s
∂
z
*
∂
s
=
∂
z
∂
(
c
*
z
)
=
c
=
∂
z
*
∂
(
c
*
z
)
=
0
再次使用 (4),我们得到
∂L
∂
z*
=c
rac{\partial L}{\partial z^*} = c
∂
z
*
∂L
=
C
。正如您所看到的,第二种方法涉及较少的计算,并且对于更快的计算更方便。
跨域功能呢? ¶
一些函数从复杂的输入映射到实际的输出,反之亦然。这些函数形成了 (4) 的特例,我们可以使用链式法则推导:
- 对于
f
:
C
→
R
f: ℂ → ℝ
f
:
C
→
R
,我们得到:
∂
L
∂
z
*
=
2
*
g
r
a
d
_
o
u
t
p>>
u
t
*
∂
s
∂
z
*
rac{\partial L}{\partial z^} = 2 * grad\ _输出 * rac{\partial s}{\partial z^{}}
∂
z
*
∂
L
=
2
*
g
r
a
d
_
o
u
tp
u
*
∂
z
*
∂
s
- 对于
f
:
R
→
C
f: ℝ → ℂ
f
:
R
→
C
,我们得到:
∂
L
∂
z
*
=
2
*
R
e
(
g
r
a
d
_
o
u
t
p
u
t
*
*
∂
s
∂
z
*
)
rac{\partial L}{\partial z^} = 2 * \mathrm{Re}(grad_output^ * * rac{\partial s}{\partial z^{}})
∂
z
*
∂
L
=
2
*
重新
(
g
r
a
d
_
o
u
tp
u
t
*
∂
z
*
∂
)
已保存tensor的挂钩 ¶
您可以通过定义一对 pack_hook
/unpack_hook
钩子来控制如何打包/解包保存的tensor。 pack_hook
函数应该采用一个tensor作为其单个参数,但可以返回任何 python 对象(例如另一个tensor、一个元组,甚至包含文件名的字符串)。 unpack_hook
函数将 pack_hook
的输出作为其单个参数,并且应该返回一个用于向后传递的tensor。 unpack_hook
返回的tensor只需与作为输入传递给 pack_hook
的tensor具有相同的内容。特别是,任何与 autograd 相关的元数据都可以忽略,因为它们将在解包过程中被覆盖。
此类对的一个示例是:
class SelfDeletingTempFile():
def __init__(self):
self.name = os.path.join(tmp_dir, str(uuid.uuid4()))
def __del__(self):
os.remove(self.name)
def pack_hook(tensor):
temp_file = SelfDeletingTempFile()
torch.save(tensor, temp_file.name)
return temp_file
def unpack_hook(temp_file):
return torch.load(temp_file.name)
请注意,unpack_hook
不应删除临时文件,因为它可能会被多次调用:只要返回的 SelfDeletingTempFile 对象处于活动状态,临时文件就应该处于活动状态。在上面的示例中,我们通过在不再需要临时文件时(删除 SelfDeletingTempFile 对象时)关闭临时文件来防止泄漏。
笔记
我们保证“pack_hook”只会被调用一次,但“unpack_hook”可以根据向后传递的需要多次调用,并且我们希望它每次都返回相同的数据。
警告
禁止对任何函数的输入执行就地操作,因为它们可能会导致意外的副作用。如果打包钩子的输入被就地修改,PyTorch 会抛出错误,但不会捕获解包钩子的输入被就地修改的情况。
为已保存的tensor注册钩子 ¶
您可以通过在“SavedTensor”对象上调用“register_hooks()”方法在已保存的tensor上注册一对钩子。这些对象作为“grad_fn”的属性公开,并以“raw_saved”前缀开头。
x = torch.randn(5, requires_grad=True)
y = x.pow(2)
y.grad_fn._raw_saved_self.register_hooks(pack_hook, unpack_hook)
一旦注册对,就会调用“pack_hook”方法。每次需要通过“y.grad_fn._saved_self”访问保存的tensor时,都会调用“unpack_hook”方法` 或在向后传递期间。
警告
如果在保存的tensor被释放后(即调用后向功能后)保留对“SavedTensor”的引用,则禁止调用其“register_hooks()”。大多数情况下,PyTorch 会抛出错误,但可能无法执行因此在某些情况下可能会出现未定义的行为。
为保存的tensor注册默认钩子 ¶
或者,您可以使用上下文管理器 saved_tensors_hooks
注册一对hooks,这将应用于在该上下文中创建的所有保存的tensor。
例子:
# Only save on disk tensors that have size >= 1000
SAVE_ON_DISK_THRESHOLD = 1000
def pack_hook(x):
if x.numel() < SAVE_ON_DISK_THRESHOLD:
return x
temp_file = SelfDeletingTempFile()
torch.save(tensor, temp_file.name)
return temp_file
def unpack_hook(tensor_or_sctf):
if isinstance(tensor_or_sctf, torch.Tensor):
return tensor_or_sctf
return torch.load(tensor_or_sctf.name)
class Model(nn.Module):
def forward(self, x):
with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
# ... compute output
output = x
return output
model = Model()
net = nn.DataParallel(model)
使用此上下文管理器定义的钩子是线程本地的。因此,以下代码将不会产生所需的效果,因为钩子不经过 DataParallel 。
# Example what NOT to do
net = nn.DataParallel(model)
with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
output = net(input)
请注意,使用这些钩子会禁用所有用于减少 Tensor 对象创建的优化。例如:
with torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x):
x = torch.randn(5, requires_grad=True)
y = x * x
如果没有钩子,x
、y.grad_fn._saved_self
和y.grad_fn._saved_other
都引用同一个tensor对象。有了钩子,PyTorch将打包和解包x 转换成两个新的tensor对象,它们与原始 x 共享相同的存储(不执行复制)。
向后钩子执行 ¶
是否会触发特定的钩子 ¶
通过 torch.Tensor.register_hook()
注册到 Tensor 的 Hook 在渐变时执行正在计算该tensor。 (请注意,这不需要执行tensor的 grad_fn。例如,如果tensor作为 inputs
参数的一部分传递给 torch.autograd.grad()
,Tensor 的 grad_fn 可能不会被执行,但该 Tensor 的钩子寄存器将始终被执行。)
通过 torch.Tensor.register_post_accumulate_grad_hook()
在该tensor的梯度累积后执行,这意味着tensor的梯度字段已设置。而通过 torch.Tensor.register_hook()
注册的钩子在梯度正在运行时运行计算,通过 torch.Tensor.register_post_accumulate_grad_hook()
注册钩子即使您调用 back(retain_graph=True) ,在非叶tensor上也会出错。
使用 torch.autograd.graph.Node.register_hook()
或 torch.autograd.graph.Node.register_prehook()
仅在其注册的节点被执行时才会被触发。
是否执行特定节点可能取决于是否使用 torch.autograd.grad()
或 torch.autograd.backward()
。具体来说,你应该注意当您在与要传递给 torch.autograd.grad()
或 torch.autograd.backward()
作为“输入”参数。
如果您使用 torch.autograd.backward()
,上述所有钩子都会无论您是否指定“inputs”参数,都会被执行。这是因为.backward() 执行所有节点,即使它们对应于指定为输入的tensor。(请注意,与作为“输入”传递的tensor相对应的此附加节点的执行通常是不必要的,但无论如何都会完成。此行为是可能会发生变化;您不应该依赖它。)
另一方面,如果您使用 torch.autograd.grad()
,则向后注册到与传递给“input”的tensor对应的节点的钩子可能不会被执行,因为除非有另一个输入依赖于该节点的梯度结果,否则这些节点不会被执行。
不同钩子的触发顺序 [¶](#the-order-in-which-the
- different-hooks-are-fired "Permalink to this header")
事情发生的顺序是:
- 执行注册到 Tensor 的钩子 2.注册到 Node 的预钩子被执行(如果 Node 被执行)。3.对于保留_grad4 的tensor,“.grad”字段会更新。节点被执行(遵守上述规则)5。对于累积了“.grad”的叶tensor,将执行 post-accumulate-grad 钩子6。注册到 Node 的 post-hooks 被执行(如果 Node 被执行)
如果在同一个 Tensor 或 Node 上注册了多个相同类型的 hook,它们将按照注册的顺序执行。后面执行的 hook 可以观察到前面的 hook 对梯度的修改。
特殊挂钩 ¶
torch.autograd.graph.register_multi_grad_hook()
是使用注册到tensor的钩子实现的。每个单独的tensor钩子都按照上面定义的tensor钩子顺序触发,并且在计算最后一个tensor梯度时调用注册的多梯度钩子。
torch.nn.modules.module.register_module_full_backward_hook()
是使用注册到 Node 的 hooks 实现的。当计算前向时,钩子被注册到与模块的输入和输出相对应的 grad_fn 。因为一个模块可能需要多个输入并返回多个输出,所以在forward之前首先将一个虚拟的自定义autograd函数应用于模块的输入,并在返回forward的输出之前应用于模块的输出,以确保这些tensor共享单个grad_fn,然后我们可以将钩子连接到上面。
就地修改tensor时tensor钩子的行为 ¶
通常注册到tensor的钩子接收相对于该tensor的输出梯度,其中tensor的值被视为向后计算时的值。
但是,如果您将钩子注册到tensor,然后就地修改该tensor,则在就地修改之前注册的钩子类似地接收输出相对于tensor的梯度,但tensor的值被视为其之前的值 -地方修改。
如果您更喜欢前一种情况的行为,则应该在对tensor进行所有就地修改后将它们注册到tensor。例如:
此外,知道在幕后,当钩子注册到tensor时,它们实际上会永久绑定到该tensor的梯度_fn,因此如果随后就地修改该tensor,即使tensor现在有一个新的 grad_fn,在就地修改之前注册的钩子将继续与旧的 grad_fn 关联,例如当 autograd 引擎在图中达到tensor的旧 grad_fn 时,它们就会触发。