自动微分包 - torch.autograd

创建于：2016 年 12 月 23 日 | 最后更新于：2025 年 6 月 12 日

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

它对现有代码的更改极少 - 您只需将需要计算梯度的 Tensor 声明为 requires_grad=True 关键字。目前，我们仅支持浮点 Tensor 类型（半精度、单精度、双精度和 bfloat16）和复数 Tensor 类型（cfloat, cdouble）的 autograd。

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

前向模式自动微分

警告

此 API 处于 Beta 阶段。尽管函数签名不太可能改变，但在我们认为其稳定之前，计划改进操作符覆盖范围。

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

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

函数式高级 API

警告

此 API 处于 Beta 阶段。尽管函数签名不太可能改变，但在我们认为其稳定之前，计划对性能进行重大改进。

本节包含基于上述基本 API 的自动微分高级 API，允许您计算雅可比矩阵、Hessian 矩阵等。

此 API 适用于仅接受张量作为输入并仅返回张量的用户提供函数。如果您的函数接受非张量参数或未设置 requires_grad 的张量，您可以使用 lambda 表达式捕获它们。例如，对于接受三个输入（一个需要计算雅可比矩阵的张量、一个应被视为常量的张量和一个布尔标志）的函数 f(input, constant, flag=flag)，您可以将其用作 functional.jacobian(lambda x: f(x, constant, flag=flag), input)。

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

局部禁用梯度计算

有关无梯度模式和推理模式以及其他可能与两者混淆的相关机制之间的差异，请参阅局部禁用梯度计算。另请参阅局部禁用梯度计算以获取可用于局部禁用梯度的函数列表。

默认梯度布局

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

如果 param.grad 最初为 None

1. 如果 param 的内存是非重叠且密集的，则 .grad 会以与 param 匹配的步长创建（从而匹配 param 的布局）。

2. 否则，.grad 将以行主连续步长创建。

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

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

4. 如果create_graph=True，backward()会将.grad替换为新的张量.grad + new grad，该操作会尝试（但不能保证）匹配原有的.grad的步长。

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

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

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

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

手动梯度布局

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

张量上的就地操作

在自动微分中支持就地操作是一个难题，我们不鼓励在大多数情况下使用它们。自动微分积极的缓冲区释放和重用使其非常高效，就地操作实际上能显著降低内存使用的情况非常少。除非您面临巨大的内存压力，否则您可能永远不需要使用它们。

就地操作正确性检查

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

变量 (已废弃)

警告

变量 API 已废弃：使用张量进行自动微分不再需要变量。自动微分自动支持requires_grad设置为True的张量。下面是关于变化的快速指南

• Variable(tensor)和Variable(tensor, requires_grad)仍然按预期工作，但它们返回的是张量而不是变量。

• 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)

张量自动微分函数

torch.Tensor.grad	此属性默认为None，并在第一次调用backward()计算self的梯度时变为张量。
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。

函数

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

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

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

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

有关如何使用此类的更多详细信息，请参见扩展 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	定义此自动微分.Function在torch.vmap()下的行为。

上下文方法混合

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

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

自定义函数工具

反向传播方法的装饰器。

function.once_differentiable

用于构建PyTorch工具的自定义Function基类

function.BackwardCFunction	此类仅用于内部自动微分工作。
function.InplaceFunction	此类仅为向后兼容性而存在。
function.NestedIOFunction	此类仅为向后兼容性而存在。

数值梯度检查

gradcheck	检查通过小有限差分计算的梯度是否与inputs中浮点或复数类型且requires_grad=True的张量相对于分析梯度一致。
gradgradcheck	检查通过小有限差分计算的梯度是否与inputs和grad_outputs中浮点或复数类型且requires_grad=True的张量相对于分析梯度一致。
GradcheckError	由gradcheck()和gradgradcheck()引发的错误。

分析器

Autograd 包含一个分析器，可让您检查模型内部不同操作符的成本——包括 CPU 和 GPU。目前实现了三种模式：仅 CPU，使用profile。基于 nvprof（同时记录 CPU 和 GPU 活动），使用emit_nvtx。以及基于 vtune 分析器，使用emit_itt。

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)

管理自动微分分析器状态并保存结果摘要的上下文管理器。

其内部机制是记录 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 的 forward 调用模块 B 的 forward，其中包含一个 aten::add 操作，则 aten::add 的模块层级是 A.B。请注意，目前此支持仅适用于 TorchScript 模型，而不适用于即时模式模型。

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

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

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

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

警告

启用内存分析或源归因会增加额外的分析器开销

警告

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

警告

由于一些 CUDA 多进程限制（参见多进程中的 CUDA），不能使用use_device = 'cuda'的分析器来对num_workers > 0的 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	在运行自动微分分析器时，为代码块/函数添加标签的上下文管理器/函数装饰器。
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)

一个上下文管理器，使每个自动微分操作都发出一个 NVTX 范围。

在nvprof下运行程序时很有用

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

不幸的是，目前无法强制 nvprof 将其收集到的数据刷新到磁盘，因此对于 CUDA 分析，必须使用此上下文管理器来注释 nvprof 跟踪并等待进程退出后才能检查它们。然后，可以使用 NVIDIA Visual Profiler (nvvp) 可视化时间线，或者torch.autograd.profiler.load_nvprof()可以加载结果以进行检查，例如在 Python REPL 中。

参数

• enabled (bool, 可选) – 设置enabled=False将使此上下文管理器变为无操作。默认值：True。

• record_shapes (bool, 可选) – 如果record_shapes=True，包裹每个自动微分操作的 nvtx 范围将附加有关该操作接收的张量参数大小的信息，格式如下：[[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...]。非张量参数将用[]表示。参数将按后端操作接收它们的顺序列出。请注意，此顺序可能与 Python 侧传递这些参数的顺序不匹配。另请注意，形状记录可能会增加 nvtx 范围创建的开销。默认值：False

示例

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

前向-后向关联

当在NVIDIA Visual Profiler中查看使用emit_nvtx创建的配置文件时，将每个后向操作与相应的前向操作相关联可能很困难。为了简化此任务，emit_nvtx会将其生成的范围附加序列号信息。

在前向传播期间，每个函数范围都用seq=<N>装饰。seq是一个计数器，每次创建并缓存一个新的反向函数对象时递增。因此，与每个前向函数范围关联的seq=<N>注释告诉您，如果此前向函数创建了反向函数对象，则反向对象将接收序列号N。在反向传播期间，包裹每个C++反向函数apply()调用的顶级范围都用stashed seq=<M>装饰。M是创建反向对象时使用的序列号。通过比较反向传播中的stashed seq数字与前向传播中的seq数字，您可以追踪是哪个前向操作创建了每个反向函数。

在反向传播过程中执行的任何函数也用seq=<N>装饰。在默认反向传播（create_graph=False）期间，此信息无关紧要，实际上，所有此类函数的N可能都只是0。只有与反向函数对象的apply()方法相关联的顶层范围才有用，作为将这些函数对象与早期前向传播关联起来的一种方式。

二次反向传播

另一方面，如果正在进行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>

Instrumentation and Tracing Technology (ITT) API 使您的应用程序能够在跨不同 Intel 工具执行期间生成和控制跟踪数据的收集。此上下文管理器用于注释 Intel(R) VTune Profiling 跟踪。借助此上下文管理器，您将能够在 Intel(R) VTune Profiler GUI 中看到带标签的范围。

参数

• enabled (bool, 可选) – 设置enabled=False将使此上下文管理器变为无操作。默认值：True。

• record_shapes (bool, 可选) – 如果record_shapes=True，则包裹每个自动微分操作的 itt 范围将附加有关该操作接收的张量参数大小的信息，格式如下：[[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...]。非张量参数将用[]表示。参数将按后端操作接收它们的顺序列出。请注意，此顺序可能与 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)

上下文管理器，根据其参数mode启用或禁用自动微分引擎的异常检测。

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

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

参数

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

• check_nan (bool) – 标志，用于在反向传播生成“nan”时是否引发错误

grad_mode.set_multithreading_enabled

上下文管理器，用于启用或禁用多线程反向传播。

自动微分图

自动微分暴露了允许检查图并在反向传播期间插入行为的方法。

torch.Tensor的grad_fn属性如果该张量是自动微分记录的操作的输出（即，grad_mode已启用且至少有一个输入需要梯度），则持有torch.autograd.graph.Node，否则为None。

graph.Node.name	返回名称。
graph.Node.metadata	返回元数据。
graph.Node.next_functions
graph.Node.register_hook	注册一个反向钩子。
graph.Node.register_prehook	注册一个反向预钩子。
graph.increment_version	更新自动微分元数据，跟踪给定张量是否被就地修改。

一些操作需要将中间结果保存在前向传播期间，以便执行反向传播。这些中间结果作为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

您还可以使用钩子定义这些保存的张量应该如何打包/解包。一个常见的应用是通过将这些中间结果保存到磁盘或 CPU 而不是留在 GPU 上来换取计算量以节省内存。如果您发现在评估期间模型适合 GPU，但在训练期间不适合，则此功能特别有用。另请参阅保存张量的钩子。

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

上下文管理器，用于为保存的张量设置一对打包/解包钩子。

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

在此上下文中，pack_hook函数将在每次操作保存张量以进行反向传播时调用（这包括使用save_for_backward()保存的中间结果，也包括由PyTorch定义的操作记录的结果）。pack_hook的输出随后存储在计算图中，而不是原始张量。

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

钩子应该具有以下签名

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的返回值不能持有对输入张量的引用。例如，使用lambda x: x.detach()而不是lambda x: x作为打包钩子。

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

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

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

使用此上下文管理器可以牺牲计算量以节省 GPU 内存使用（例如，当您的模型在训练期间不适合 GPU 内存时）。

参数

pin_memory (bool) – 如果为True，张量在打包期间将保存到 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)

上下文管理器，用于禁用保存张量的默认钩子功能。

如果您正在创建不与保存张量默认钩子一起使用的功能，则很有用。

参数

error_message (str) – 当保存张量默认钩子被禁用时，如果仍使用它们，则会引发一个带有此错误消息的 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中每个张量的梯度后调用。如果张量在tensors中但不是图的一部分，或者如果张量对于计算当前.backward()或.grad()调用中指定的任何inputs的梯度不是必需的，则此张量将被忽略，并且钩子将不会等待其梯度被计算。

在所有未被忽略的张量的梯度计算完毕后，fn将与这些梯度一起被调用。None将被传递给那些没有计算梯度的张量。

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

钩子不应修改其参数。

此函数返回一个带有方法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

在此上下文管理器下，允许对为反向传播保存的张量进行修改。

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

为确保正确行为，前向和反向传播应在同一上下文管理器下运行。

返回

一个存储此上下文管理器所管理状态的 _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)

表示自动微分图中给定梯度边的对象。

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

torch.autograd.graph.get_gradient_edge(tensor)

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

特别是，它等价于调用g = autograd.grad(loss, input)和g = autograd.grad(loss, get_gradient_edge(input))。

返回类型

GradientEdge