自动微分包 - torch.autograd

创建日期：2016年12月23日 | 最后更新日期：2025年11月01日

torch.autograd 提供了实现任意标量值函数自动微分的类和函数。

它对现有代码的修改极小——你只需要使用 requires_grad=True 关键字声明需要计算梯度的 Tensor。目前，我们仅支持对浮点 Tensor 类型（half, float, double 和 bfloat16）以及复数 Tensor 类型（cfloat, cdouble）进行自动微分。

backward	计算给定张量相对于图叶子节点的梯度和。
grad	计算并返回输出相对于输入的梯度和。

前向模式自动微分 (Forward-mode Automatic Differentiation)

警告

此 API 处于测试阶段。尽管函数签名极不可能更改，但我们计划在将其视为稳定之前提升算子覆盖范围。

有关如何使用此 API 的详细步骤，请参阅前向模式 AD 教程。

forward_ad.dual_level	前向 AD 的上下文管理器，所有的前向 AD 计算必须在 dual_level 上下文中进行。
forward_ad.make_dual	将张量值与其切线关联，从而为前向 AD 梯度计算创建一个“对偶张量”。
forward_ad.unpack_dual	解包一个“对偶张量”，以获取其原始 Tensor 值及其前向 AD 梯度。
forward_ad.enter_dual_level	进入一个新的前向梯度级别。
forward_ad.exit_dual_level	退出一个前向梯度级别。
forward_ad.UnpackedDualTensor	unpack_dual() 返回的具名元组，包含对偶张量的原始分量和切线分量。

函数式高阶 API (Functional higher level API)

警告

此 API 处于测试阶段。尽管函数签名极不可能更改，但我们计划在将其视为稳定之前对性能进行重大改进。

本节包含了 autograd 的高阶 API，它建立在上述基础 API 之上，允许你计算雅可比矩阵、海森矩阵等。

此 API 适用于用户提供的仅以 Tensor 为输入并仅返回 Tensor 的函数。如果你的函数接收非 Tensor 参数，或者接收未设置 requires_grad 的 Tensor，可以使用 lambda 来捕获它们。例如，对于一个函数 f，它接收三个输入：一个我们要计算雅可比矩阵的 Tensor，另一个应被视为常量的张量，以及一个布尔标志 f(input, constant, flag=flag)，你可以这样使用它：functional.jacobian(lambda x: f(x, constant, flag=flag), input)。

functional.jacobian	计算给定函数的雅可比矩阵。
functional.hessian	计算给定标量函数的海森矩阵。
functional.vjp	计算向量 v 与给定函数在输入点处的雅可比矩阵的点积。
functional.jvp	计算给定函数在输入点处的雅可比矩阵与向量 v 的点积。
functional.vhp	计算向量 v 与给定标量函数在指定点处的海森矩阵的点积。
functional.hvp	计算标量函数的海森矩阵与向量 v 在指定点处的点积。

局部禁用梯度计算

请参阅 局部禁用梯度计算 以获取有关 no-grad（无梯度）和推断模式之间的区别，以及可能与两者混淆的其他相关机制的更多信息。另请参阅 局部禁用梯度计算 以获取可用于局部禁用梯度的函数列表。

默认梯度布局 (Default gradient layouts)

当非稀疏 param 在 torch.autograd.backward() 或 torch.Tensor.backward() 期间接收非稀疏梯度时，param.grad 按如下方式累加。

如果 param.grad 最初为 None

1. 如果 param 的内存是非重叠且稠密的，则创建 .grad 时，其步长（strides）将与 param 匹配（因此与 param 的布局匹配）。

2. 否则，.grad 将以行优先连续步长（rowmajor-contiguous strides）创建。

如果 param 已经具有非稀疏的 .grad 属性

3. 如果 create_graph=False，backward() 将原地累加到 .grad 中，这会保留其步长。

4. 如果 create_graph=True，backward() 会用一个新张量 .grad + new grad 替换 .grad，它会尝试（但不保证）匹配原有的 .grad 的步长。

为了获得最佳性能，建议采用默认行为（即在第一次 backward() 之前让 .grads 为 None，以便按照 1 或 2 创建它们的布局，并根据 3 或 4 随时间保留它们）。调用 model.zero_grad() 或 optimizer.zero_grad() 不会影响 .grad 的布局。

事实上，在每个累加阶段之前将所有 .grads 重置为 None，例如：

for iterations...
    ...
    for param in model.parameters():
        param.grad = None
    loss.backward()

即每次都根据 1 或 2 重新创建它们，是 model.zero_grad() 或 optimizer.zero_grad() 的一种有效替代方案，这可能会提升某些网络的性能。

手动梯度布局 (Manual gradient layouts)

如果你需要手动控制 .grad 的步长，请在第一次 backward() 之前将 param.grad 赋值为一个具有所需步长的全零张量，且不要将其重置为 None。只要 create_graph=False，3 保证你的布局会被保留。4 指出即使 create_graph=True，你的布局也很可能会被保留。

对张量的原地操作 (In-place operations on Tensors)

在 autograd 中支持原地操作是一个难题，我们在大多数情况下不鼓励使用它们。Autograd 积极的缓冲区释放和复用使其非常高效，实际上很少有情况原地操作能显著降低内存使用量。除非你是在高内存压力下操作，否则你可能永远不需要使用它们。

原地操作正确性检查 (In-place correctness checks)

所有的 Tensor 都会跟踪应用在其上的原地操作。如果实现检测到某个张量在函数中被保存用于反向传播，但随后被原地修改了，那么在反向传播开始时会报错。这确保了如果你在使用原地函数且没有看到任何错误，那么你可以确信计算出的梯度是正确的。

Variable (已弃用)

警告

Variable API 已被弃用：使用 autograd 和张量时不再需要 Variables。Autograd 自动支持 requires_grad 设置为 True 的张量。以下是关于变化情况的简要指南：

• Variable(tensor) 和 Variable(tensor, requires_grad) 仍然可以按预期工作，但它们返回的是 Tensors 而不是 Variables。

• var.data 与 tensor.data 是同一回事。

• 诸如 var.backward(), var.detach(), var.register_hook() 之类的方法现在可以直接用于具有相同方法名的张量。

此外，现在可以使用工厂方法创建 requires_grad=True 的张量，例如 torch.randn(), torch.zeros(), torch.ones() 以及其他如下所示的方法：

autograd_tensor = torch.randn((2, 3, 4), requires_grad=True)

张量 autograd 函数 (Tensor autograd functions)

torch.Tensor.grad	此属性默认值为 None，并在第一次调用 backward() 计算 self 的梯度时变为一个 Tensor。
torch.Tensor.requires_grad	如果需要为此张量计算梯度，则为 True，否则为 False。
torch.Tensor.is_leaf	所有 requires_grad 为 False 的张量按照约定都将是叶子张量。
torch.Tensor.backward([gradient, ...])	计算当前张量相对于图叶的梯度。
torch.Tensor.detach	返回一个从当前图分离的新张量。
torch.Tensor.detach_	将张量从创建它的图中分离出来，使其成为叶子节点。
torch.Tensor.register_hook(hook)	注册一个反向钩子。
torch.Tensor.register_post_accumulate_grad_hook(hook)	注册一个在梯度累积后运行的反向钩子。
torch.Tensor.retain_grad()	使该张量能够在 backward() 期间填充其 grad。

Function

class torch.autograd.Function(*args, **kwargs)

创建自定义 autograd.Function 的基类。

要创建自定义 autograd.Function，请继承此类并实现 forward() 和 backward() 静态方法。然后，要在前向传播中使用你的自定义操作，请调用类方法 apply。不要直接调用 forward()。

为了确保正确性和最佳性能，请确保在 ctx 上调用了正确的方法，并使用 torch.autograd.gradcheck() 验证你的 backward 函数。

有关如何使用该类的更多详细信息，请参阅 扩展 torch.autograd。

示例

>>> class Exp(Function):
>>>     @staticmethod
>>>     def forward(ctx, i):
>>>         result = i.exp()
>>>         ctx.save_for_backward(result)
>>>         return result
>>>
>>>     @staticmethod
>>>     def backward(ctx, grad_output):
>>>         result, = ctx.saved_tensors
>>>         return grad_output * result
>>>
>>> # Use it by calling the apply method:
>>> output = Exp.apply(input)

Function.forward	定义自定义自动微分函数的前向传播。
Function.backward	定义使用反向模式自动微分来区分操作的公式。
Function.jvp	定义使用前向模式自动微分来区分操作的公式。
Function.vmap	定义此 autograd.Function 在 torch.vmap() 下的行为。

上下文方法 Mixins (Context method mixins)

创建新的 Function 时，以下方法可用于 ctx。

function.FunctionCtx.mark_dirty	将给定张量标记为在就地操作中已修改。
function.FunctionCtx.mark_non_differentiable	将输出标记为不可微分。
function.FunctionCtx.save_for_backward	保存给定的张量以供将来调用 backward()。
function.FunctionCtx.set_materialize_grads	设置是否具体化（materialize）梯度张量。

自定义 Function 实用程序 (Custom Function utilities)

用于 backward 方法的装饰器。

function.once_differentiable

用于构建 PyTorch 实用程序的自定义 Function 基类

function.BackwardCFunction	此类用于内部 autograd 工作。
function.InplaceFunction	此类仅出于向后兼容性原因存在。
function.NestedIOFunction	此类仅出于向后兼容性原因存在。

数值梯度检查 (Numerical gradient checking)

gradcheck	针对 inputs 中具有浮点或复数类型且 requires_grad=True 的张量，通过微小的有限差分来检查通过解析计算得出的梯度。
gradgradcheck	针对 inputs 和 grad_outputs 中具有浮点或复数类型且 requires_grad=True 的张量，通过微小的有限差分来检查梯度计算出的梯度与解析梯度。
GradcheckError	gradcheck() 和 gradgradcheck() 引发的错误。

性能分析器 (Profiler)

Autograd 包含一个性能分析器，让你能够检查模型中不同算子的开销（无论是在 CPU 还是 GPU 上）。目前实现了三种模式：使用 profile 进行仅 CPU 分析；使用 emit_nvtx 进行基于 nvprof 的分析（记录 CPU 和 GPU 活动）；以及使用 emit_itt 的基于 vtune 的分析。

class torch.autograd.profiler.profile(enabled=True, *, use_cuda=False, use_device=None, record_shapes=False, with_flops=False, profile_memory=False, with_stack=False, with_modules=False, use_kineto=False, use_cpu=True, experimental_config=None, acc_events=False, custom_trace_id_callback=None, post_processing_timeout_s=None)

管理 autograd 性能分析器状态并保存结果摘要的上下文管理器。

注意

这是后端，大多数用户应该使用 torch.profiler。

底层它只是记录在 C++ 中执行的函数事件，并将这些事件公开给 Python。你可以将任何代码包裹在其中，它只会报告 PyTorch 函数的运行时间。注意：分析器是线程局部的，并会自动传播到异步任务中。

参数:

• enabled (bool, 可选) – 将其设置为 False 使此上下文管理器无效。

• use_cuda (bool, 可选) – 使用 cudaEvent API 同时启用 CUDA 事件计时。（即将被弃用）

• use_device (str, 可选) – 启用设备事件计时。使用 cuda 时，每个张量操作会增加大约 4us 的开销。有效的设备选项包括 ‘cuda’, ‘xpu’, ‘mtia’ 和 ‘privateuseone’。

• record_shapes (bool, 可选) – 如果设置了形状记录，将收集有关输入维度的信息。这允许查看底层使用了哪些维度，并使用 prof.key_averages(group_by_input_shape=True) 按维度进行进一步分组。请注意，形状记录可能会使你的分析数据产生偏差。建议分别进行形状记录和非形状记录的运行以验证计时。偏差对于最底层的事件（在嵌套函数调用的情况下）通常是可以忽略不计的。但对于高级函数，由于形状收集，总的 CPU 自用时间可能会被人为夸大。

• with_flops (bool, 可选) – 如果设置了 with_flops，分析器将使用算子的输入形状估计 FLOPs（浮点运算）值。这允许估算硬件性能。目前，此选项仅适用于矩阵乘法和 2D 卷积算子。

• profile_memory (bool, 可选) – 跟踪张量内存分配/释放。

• with_stack (bool, 可选) – 记录算子的源信息（文件和行号）。

• with_modules (bool) – 记录对应于算子调用栈的模块层次结构（包括函数名）。例如：如果模块 A 的前向调用了模块 B 的前向，其中包含 aten::add 算子，那么 aten::add 的模块层次结构就是 A.B。请注意，目前此支持仅存在于 TorchScript 模型，而不适用于 Eager 模式模型。

• use_kineto (bool, 可选) – 实验性功能，启用 Kineto 分析器进行性能分析。

• use_cpu (bool, 可选) – 分析 CPU 事件；设置为 False 需要 use_kineto=True，可用于降低仅 GPU 分析时的开销。

• experimental_config (_ExperimentalConfig) – 一组由 Kineto 等分析器库使用的实验性选项。请注意，不保证向后兼容性。

• acc_events (bool) – 启用跨多个分析周期累积 FunctionEvents。

• post_processing_timeout_s (float) – 性能分析结果后处理的可选超时时间（秒）。在此上下文中，后处理发生在分析本身完成后。如果指定，事件解析将在该持续时间后停止并返回部分结果。这对于处理可能需要太长时间处理的大规模跟踪数据非常有用。

警告

启用内存分析或源代码归属会产生额外的分析器开销。

警告

此上下文管理器不应递归调用，即不允许嵌套实例。

警告

由于某些 CUDA 多进程限制（参见 CUDA 中的多进程），无法在 num_workers > 0 的情况下使用 use_device = 'cuda' 来对 DataLoaders 进行基准测试。如果你希望对数据加载进行基准测试，请使用 use_device = None 或 num_workers = 0。

示例

>>> x = torch.randn((1, 1), requires_grad=True)
>>> with torch.autograd.profiler.profile() as prof:
>>>     for _ in range(100):  # any normal python code, really!
>>>         y = x ** 2
>>>         y.backward()
>>> # NOTE: some columns were removed for brevity
>>> print(prof.key_averages().table(sort_by="self_cpu_time_total"))
-----------------------------------  ---------------  ---------------  ---------------
Name                                 Self CPU total   CPU time avg     Number of Calls
-----------------------------------  ---------------  ---------------  ---------------
mul                                  32.048ms         32.048ms         200
pow                                  27.041ms         27.041ms         200
PowBackward0                         9.727ms          55.483ms         100
torch::autograd::AccumulateGrad      9.148ms          9.148ms          100
torch::autograd::GraphRoot           691.816us        691.816us        100
-----------------------------------  ---------------  ---------------  ---------------

profiler.profile.export_chrome_trace	将 EventList 导出为 Chrome 跟踪工具文件。
profiler.profile.key_averages	计算所有函数事件在其键上的平均值。
profiler.profile.self_cpu_time_total	返回在 CPU 上花费的总时间。
profiler.profile.total_average	计算所有事件的聚合统计信息。
profiler.parse_nvprof_trace
profiler.EnforceUnique	如果某个键出现多次，则引发错误。
profiler.KinetoStepTracker	提供全局递增步骤计数的抽象。
profiler.record_function	上下文管理器/函数装饰器，在运行 autograd 性能分析器时为代码块/函数添加标签。
profiler_util.Interval
profiler_util.Kernel
profiler_util.MemRecordsAcc	用于在区间内访问 mem_records 的加速结构。
profiler_util.StringTable

class torch.autograd.profiler.emit_nvtx(enabled=True, record_shapes=False)

使每个 autograd 操作发射 NVTX 范围的上下文管理器。

当在 nvprof 下运行程序时非常有用。

nvprof --profile-from-start off -o trace_name.prof -- <regular command here>

遗憾的是，无法强制 nvprof 将其收集的数据刷新到磁盘，因此在进行 CUDA 性能分析时，必须使用此上下文管理器来标注 nvprof 跟踪信息，并等待进程结束后才能进行检查。之后，既可以使用 NVIDIA 可视化分析器 (nvvp) 来可视化时间轴，也可以使用 torch.autograd.profiler.load_nvprof() 加载结果，例如在 Python REPL 中进行查看。

参数:

• enabled (bool, optional) – 设置 enabled=False 会使该上下文管理器失效（no-op）。默认值：True。

• record_shapes (bool, optional) – 如果 record_shapes=True，包裹每个自动求导（autograd）算子的 nvtx 范围将追加该算子接收到的 Tensor 参数尺寸信息，格式如下：[[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...] 非 Tensor 参数将表示为 []。参数将按照它们被后端算子接收的顺序排列。请注意，此顺序可能与 Python 端传递这些参数的顺序不一致。另请注意，形状记录可能会增加创建 nvtx 范围的开销。默认值：False

示例

>>> with torch.cuda.profiler.profile():
...     model(x)  # Warmup CUDA memory allocator and profiler
...     with torch.autograd.profiler.emit_nvtx():
...         model(x)

前向-后向关联

在使用 emit_nvtx 创建的 profile 文件并在 Nvidia Visual Profiler 中查看时，将每个后向传播算子与相应的前向传播算子关联起来可能很困难。为了简化此任务，emit_nvtx 会在生成的范围中追加序列号信息。

在前向传播过程中，每个函数范围都会被标记为 seq=<N>。seq 是一个运行计数器，每当创建一个新的后向 Function 对象并为了后向传播而存入（stash）时，它就会增加。因此，与每个前向函数范围相关联的 seq=<N> 注解告诉您：如果该前向函数创建了一个后向 Function 对象，那么该后向对象将获得序列号 N。在后向传播过程中，包裹每个 C++ 后向 Function 的 apply() 调用的顶层范围被标记为 stashed seq=<M>。M 是后向对象创建时的序列号。通过比较后向传播中的 stashed seq 号码与前向传播中的 seq 号码，您可以追踪是哪个前向算子创建了每个后向 Function。

在后向传播期间执行的任何函数也会被标记为 seq=<N>。在默认的后向传播（create_graph=False）期间，此信息无关紧要，事实上，所有这些函数的 N 可能都只是 0。只有与后向 Function 对象的 apply() 方法关联的顶层范围是有用的，作为将这些 Function 对象与较早的前向传播关联起来的一种方式。

二阶后向传播（Double-backward）

另一方面，如果正在进行 create_graph=True 的后向传播（即，如果您正在为二阶后向传播做准备），则后向传播期间每个函数的执行都会被赋予一个非零且有用的 seq=<N>。这些函数本身也可能创建需要在二阶后向传播期间稍后执行的 Function 对象，就像前向传播中的原始函数所做的那样。后向传播和二阶后向传播之间的关系在概念上与前向传播和后向传播之间的关系相同：函数仍然发出带有当前序列号标记的范围，它们创建的 Function 对象仍然存入这些序列号，并且在最终的二阶后向传播期间，Function 对象的 apply() 范围仍然被标记为 stashed seq 号码，这些号码可以与后向传播中的 seq 号码进行比较。

class torch.autograd.profiler.emit_itt(enabled=True, record_shapes=False)

使每个自动求导操作发出 ITT 范围的上下文管理器。

它在 Intel(R) VTune Profiler 下运行程序时非常有用。

vtune <--vtune-flags> <regular command here>

仪表化和跟踪技术 (ITT) API 使您的应用程序能够生成并控制在跨不同 Intel 工具执行期间的跟踪数据收集。此上下文管理器用于对 Intel(R) VTune Profiling 跟踪进行标注。借助此上下文管理器，您将能够在 Intel(R) VTune Profiler GUI 中看到标记的范围。

参数:

• enabled (bool, optional) – 设置 enabled=False 会使该上下文管理器失效（no-op）。默认值：True。

• record_shapes (bool, optional) – 如果 record_shapes=True，包裹每个自动求导算子的 itt 范围将追加该算子接收到的 Tensor 参数的尺寸信息，格式如下：[[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...] 非 Tensor 参数将表示为 []。参数将按照它们被后端算子接收的顺序排列。请注意，此顺序可能与 Python 端传递这些参数的顺序不一致。另请注意，形状记录可能会增加创建 itt 范围的开销。默认值：False

示例

>>> with torch.autograd.profiler.emit_itt():
...     model(x)

profiler.load_nvprof

打开一个 nvprof 跟踪文件并解析自动求导注解。

调试与异常检测

class torch.autograd.detect_anomaly(check_nan=True)

为自动求导引擎启用异常检测的上下文管理器。

它执行两项操作

• 在启用检测的情况下运行前向传播，将允许后向传播打印创建导致错误后向函数的那个前向操作的跟踪信息。

• 如果 check_nan 为 True，任何产生 “nan” 值的后向计算都将引发错误。默认值为 True。

警告

此模式仅应在调试时启用，因为不同的测试会拖慢程序的执行速度。

示例

>>> import torch
>>> from torch import autograd
>>> class MyFunc(autograd.Function):
...     @staticmethod
...     def forward(ctx, inp):
...         return inp.clone()
...
...     @staticmethod
...     def backward(ctx, gO):
...         # Error during the backward pass
...         raise RuntimeError("Some error in backward")
...         return gO.clone()
>>> def run_fn(a):
...     out = MyFunc.apply(a)
...     return out.sum()
>>> inp = torch.rand(10, 10, requires_grad=True)
>>> out = run_fn(inp)
>>> out.backward()
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
      File "/your/pytorch/install/torch/_tensor.py", line 93, in backward
        torch.autograd.backward(self, gradient, retain_graph, create_graph)
      File "/your/pytorch/install/torch/autograd/__init__.py", line 90, in backward
        allow_unreachable=True)  # allow_unreachable flag
      File "/your/pytorch/install/torch/autograd/function.py", line 76, in apply
        return self._forward_cls.backward(self, *args)
      File "<stdin>", line 8, in backward
    RuntimeError: Some error in backward
>>> with autograd.detect_anomaly():
...     inp = torch.rand(10, 10, requires_grad=True)
...     out = run_fn(inp)
...     out.backward()
    Traceback of forward call that caused the error:
      File "tmp.py", line 53, in <module>
        out = run_fn(inp)
      File "tmp.py", line 44, in run_fn
        out = MyFunc.apply(a)
    Traceback (most recent call last):
      File "<stdin>", line 4, in <module>
      File "/your/pytorch/install/torch/_tensor.py", line 93, in backward
        torch.autograd.backward(self, gradient, retain_graph, create_graph)
      File "/your/pytorch/install/torch/autograd/__init__.py", line 90, in backward
        allow_unreachable=True)  # allow_unreachable flag
      File "/your/pytorch/install/torch/autograd/function.py", line 76, in apply
        return self._forward_cls.backward(self, *args)
      File "<stdin>", line 8, in backward
    RuntimeError: Some error in backward

class torch.autograd.set_detect_anomaly(mode, check_nan=True)

用于设置自动求导引擎的异常检测开启或关闭的上下文管理器。

set_detect_anomaly 将根据其参数 mode 启用或禁用自动求导异常检测。它可以用作上下文管理器或函数。

有关异常检测行为的详细信息，请参见上方的 detect_anomaly。

参数:

• mode (bool) – 标记是否启用异常检测（True），或禁用（False）。

• check_nan (bool) – 标记后向传播生成 “nan” 时是否引发错误

grad_mode.set_multithreading_enabled

启用或禁用多线程后向传播的上下文管理器。

自动求导图（Autograd graph）

自动求导提供了一些方法，允许检查计算图并在后向传播期间干预其行为。

如果 Tensor 是由自动求导记录的操作的输出（即 grad_mode 处于启用状态，且至少有一个输入需要梯度），则 torch.Tensor 的 grad_fn 属性会持有一个 torch.autograd.graph.Node，否则为 None。

graph.Node.name	返回名称。
graph.Node.metadata	返回元数据。
graph.Node.next_functions
graph.Node.register_hook	注册一个后向钩子（backward hook）。
graph.Node.register_prehook	注册一个后向预钩子（backward pre-hook）。
graph.increment_version	更新自动求导元数据，用于跟踪给定的 Tensor 是否进行了原地（in-place）修改。

某些操作需要保存前向传播期间的中间结果，以便执行后向传播。这些中间结果作为属性保存在 grad_fn 上，并且可以被访问。例如

>>> a = torch.tensor([0., 0., 0.], requires_grad=True)
>>> b = a.exp()
>>> print(isinstance(b.grad_fn, torch.autograd.graph.Node))
True
>>> print(dir(b.grad_fn))
['__call__', '__class__', '__delattr__', '__dir__', '__doc__', '__eq__', '__format__', '__ge__', '__getattribute__', '__gt__', '__hash__', '__init__', '__init_subclass__', '__le__', '__lt__', '__ne__', '__new__', '__reduce__', '__reduce_ex__', '__repr__', '__setattr__', '__sizeof__', '__str__', '__subclasshook__', '_raw_saved_result', '_register_hook_dict', '_saved_result', 'metadata', 'name', 'next_functions', 'register_hook', 'register_prehook', 'requires_grad']
>>> print(torch.allclose(b.grad_fn._saved_result, b))
True

您还可以使用钩子定义这些保存的 Tensor 应如何打包/解包。一个常见的应用是通过将这些中间结果保存到磁盘或 CPU 而不是保留在 GPU 上，从而以计算换取内存。如果您发现模型在评估时能适应 GPU，但在训练时不行，这尤其有用。另请参见 保存 Tensor 的钩子。

class torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook)

为保存的 Tensor 设置一对打包/解包钩子的上下文管理器。

使用此上下文管理器来定义操作的中间结果在保存前应如何打包，以及在检索时应如何解包。

在此上下文中，每当操作保存用于后向传播的 Tensor 时，都会调用 pack_hook 函数（这包括使用 save_for_backward() 保存的中间结果，也包括 PyTorch 定义的操作记录的中间结果）。pack_hook 的输出随后会被存储在计算图中，而不是原始的 Tensor。

当需要访问保存的 Tensor 时，即执行 torch.Tensor.backward() 或 torch.autograd.grad() 时，会调用 unpack_hook。它将 pack_hook 返回的 *打包* 对象作为参数，并应返回一个内容与原始 Tensor（传递给相应 pack_hook 的输入）相同的 Tensor。

钩子应具有以下签名

pack_hook(tensor: Tensor) -> Any

unpack_hook(Any) -> Tensor

其中 pack_hook 的返回值必须是 unpack_hook 的有效输入。

通常，您希望 unpack_hook(pack_hook(t)) 在值、大小、数据类型和设备方面等于 t。

示例

>>> def pack_hook(x):
...     print("Packing", x)
...     return x.detach()
>>>
>>> def unpack_hook(x):
...     print("Unpacking", x)
...     return x
>>>
>>> a = torch.ones(5, requires_grad=True)
>>> b = torch.ones(5, requires_grad=True) * 2
>>> with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
...     y = a * b
Packing tensor([1., 1., 1., 1., 1.], requires_grad=True)
Packing tensor([2., 2., 2., 2., 2.], grad_fn=<MulBackward0>)
>>> y.sum().backward()
Unpacking tensor([1., 1., 1., 1., 1.], requires_grad=True)
Unpacking tensor([2., 2., 2., 2., 2.], grad_fn=<MulBackward0>)

警告

对任一钩子的输入执行原地操作可能会导致未定义的行为。

警告

一次只允许存在一对钩子。当递归嵌套此上下文管理器时，仅应用最内层的钩子对。

警告

为了避免引用循环，pack_hook 的返回值不能持有对输入 Tensor 的引用。例如，使用 lambda x: x.detach() 而不是 lambda x: x 作为打包钩子。

class torch.autograd.graph.save_on_cpu(pin_memory=False, device_type='cuda')

在该上下文管理器下，前向传播保存的 Tensor 将存储在 CPU 上，然后在后向传播时检索。

在执行此上下文管理器内的操作时，前向传播期间保存在计算图中的中间结果将被移动到 CPU，然后在需要后向传播时复制回原始设备。如果计算图已经在 CPU 上，则不执行 Tensor 复制。

使用此上下文管理器以计算换取 GPU 内存使用（例如，当您的模型在训练期间无法放入 GPU 内存时）。

参数:

pin_memory (bool) – 如果为 True，Tensor 将在打包期间保存到 CPU 固定内存中，并在解包期间异步复制到 GPU。默认值为 False。另请参见 使用固定内存缓冲区。

示例

>>> a = torch.randn(5, requires_grad=True, device="cuda")
>>> b = torch.randn(5, requires_grad=True, device="cuda")
>>> c = torch.randn(5, requires_grad=True, device="cuda")
>>>
>>> def f(a, b, c):
...     prod_1 = a * b           # a and b are saved on GPU
...     with torch.autograd.graph.save_on_cpu():
...         prod_2 = prod_1 * c  # prod_1 and c are saved on CPU
...     y = prod_2 * a           # prod_2 and a are saved on GPU
...     return y
>>>
>>> y = f(a, b, c)
>>> del a, b, c  # for illustration only
>>> # the content of a, b, and prod_2 are still alive on GPU
>>> # the content of prod_1 and c only live on CPU
>>> y.sum().backward()  # all CPU tensors are moved back to GPU, for backward
>>> # all intermediary tensors are released (deleted) after the call to backward

class torch.autograd.graph.disable_saved_tensors_hooks(error_message)

禁用保存的 Tensor 默认钩子功能的上下文管理器。

如果您正在创建一个不适用于保存 Tensor 默认钩子的功能，这非常有用。

参数:

error_message (str) – 当在禁用保存 Tensor 默认钩子的情况下尝试使用它们时，将引发带有此错误消息的 RuntimeError。

返回类型:

Generator[None, None, None]

示例

>>> message = "saved tensors default hooks are disabled"
>>> with torch.autograd.graph.disable_saved_tensors_hooks(message):
...     # Raises RuntimeError: saved tensors default hooks are disabled
...     with torch.autograd.graph.save_on_cpu():
...         pass

class torch.autograd.graph.register_multi_grad_hook(tensors, fn, *, mode='all')

注册一个多梯度后向钩子。

支持两种模式："all" 和 "any"。

在 "all" 模式下，钩子将在计算出 tensors 中每个 Tensor 的梯度后被调用。如果一个 Tensor 在 tensors 中但不是计算图的一部分，或者如果计算当前 .backward() 或 .grad() 调用指定的任何 inputs 不需要该 Tensor 的梯度，则该 Tensor 将被忽略，钩子不会等待其梯度被计算。

在每个未被忽略的 Tensor 的梯度计算完成后，将使用这些梯度调用 fn。对于未计算梯度的 Tensor，将传入 None。

在 "any" 模式下，钩子将在计算出 tensors 中任意一个 Tensor 的第一个梯度后被调用。钩子将以该梯度作为参数被调用。

钩子不应修改其参数。

此函数返回一个包含 handle.remove() 方法的句柄，用于移除钩子。

注意

请参见 后向钩子执行，了解有关此钩子何时执行以及其执行顺序相对于其他钩子如何排序的更多信息。

示例

>>> import torch
>>>
>>> a = torch.rand(2, 3, requires_grad=True)
>>> b = torch.rand(2, 3, requires_grad=True)
>>> c = a * b
>>> d = a * b
>>>
>>> def fn(grads):
...     print([g is not None for g in grads])
...
>>> torch.autograd.graph.register_multi_grad_hook((a, b, c, d), fn)
>>>
>>> c.sum().backward(retain_graph=True)
[True, True, True, False]
>>> c.sum().backward(inputs=(a,), retain_graph=True)
[True, False, True, False]
>>>

返回类型:

RemovableHandle

class torch.autograd.graph.allow_mutation_on_saved_tensors

允许对为后向传播保存的 Tensor 进行修改的上下文管理器。

在此上下文管理器下，为后向传播保存的 Tensor 在修改时会被克隆，因此原始版本仍然可以在后向传播期间使用。通常，修改为后向传播保存的 Tensor 将导致在后向传播期间使用它时引发错误。

为了确保正确的行为，前向传播和后向传播都应在同一个上下文管理器下运行。

返回:

一个 _AllowMutationOnSavedContext 对象，用于存储由此上下文管理器管理的状态。该对象对于调试目的非常有用。上下文管理器管理的状态会在退出时自动清除。

返回类型:

Generator[_AllowMutationOnSavedContext, None, None]

示例

>>> import torch
>>> with torch.autograd.graph.allow_mutation_on_saved_tensors():
...     # forward
...     a = torch.ones(2, 3, requires_grad=True)
...     b = a.clone()
...     out = (b**2).sum()
...     b.sin_()
...     # backward
...     out.sum().backward()
...
tensor([[0.8415, 0.8415, 0.8415],
        [0.8415, 0.8415, 0.8415]], grad_fn=<SinBackward0>)

class torch.autograd.graph.GradientEdge(node, output_nr, ownership_token=None)

表示自动求导图中给定梯度边的对象。

要获取将计算给定 Tensor 梯度的梯度边，可以执行 edge = autograd.graph.get_gradient_edge(tensor)。

torch.autograd.graph.get_gradient_edge(tensor)

获取用于计算给定 Tensor 梯度的梯度边。

特别地，调用 g = autograd.grad(loss, input) 与 g = autograd.grad(loss, get_gradient_edge(input)) 是等价的。

返回类型:

GradientEdge

torch.autograd.graph.set_warn_on_accumulate_grad_stream_mismatch(enabled)

是否在 AccumulateGrad 节点的流与产生传入梯度的节点的流不匹配时发出警告。