自动微分包 - torch.autograd ¶

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

它需要最小化修改现有代码 - 您只需使用 requires_grad=True 关键字声明 Tensor ，对于其梯度需要计算的变量。目前，我们仅支持浮点 Tensor 类型（half、float、double 和 bfloat16）和复数 Tensor 类型（cfloat、cdouble）。

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

前向模式自动微分

警告

此 API 处于测试版。尽管函数签名非常不可能更改，但在我们将其视为稳定版之前，计划改进操作覆盖范围。

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

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() 返回的 Namedtuple，包含对偶张量的原点和切线分量。

函数高级 API ¶

警告

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

本节包含高级 API，用于自动微分，它基于上面的基本 API，并允许您计算雅可比矩阵、海森矩阵等。

此 API 与用户提供的函数一起工作，这些函数只接受张量作为输入并返回张量。如果您的函数接受其他不是张量或没有设置 requires_grad 属性的张量作为参数，您可以使用 lambda 表达式来捕获它们。例如，对于需要三个输入的函数 f ，一个我们想要计算雅可比矩阵的张量，另一个应被视为常量的张量，以及一个布尔标志 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 在指定点的点积。

局部禁用梯度计算

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

默认梯度布局

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

如果 param.grad 最初是 None ：

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

2. 否则， .grad 将使用 rowmajor-contiguous 步长创建。

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

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

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

默认行为（在第一次 .grad 之前让 None 保持 backward() ，以便它们的布局根据 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 ，您的布局也很可能得到保留。

张量原地操作 ¶

在 autograd 中支持原地操作是一项艰巨的任务，我们通常不建议使用它们。Autograd 的积极缓冲区释放和重用使其非常高效，而且很少情况下原地操作能显著降低内存使用量。除非你处于严重的内存压力之下，否则你可能永远不需要使用它们。

原地操作正确性检查 ¶

所有 Tensor 都记录了对其应用的原地操作，如果实现检测到某个张量在一个函数中为反向传播保存了，但之后被原地修改了，一旦开始反向传播，将会引发错误。这确保了如果你使用原地函数且没有看到任何错误，你可以确信计算出的梯度是正确的。

变量（已弃用）¶

警告

变量 API 已被弃用：使用自动微分与张量不再需要变量。自动微分自动支持将 requires_grad 设置为 True 的张量。以下是一个快速指南，说明有哪些变化：

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

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

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

此外，现在可以使用工厂方法 torch.randn() 、 torch.zeros() 、 torch.ones() 等创建张量，如下所示：

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

张量自动微分函数

torch.Tensor.grad	此属性默认为 None ，并在第一次调用 backward() 计算 self 的梯度时变为 Tensor。
torch.Tensor.requires_grad	如果需要为此张量计算梯度，则使用 True ；否则使用 False 。
torch.Tensor.is_leaf	所有具有 requires_grad 且为 False 的 Tensor 都将按照惯例成为叶子 Tensor。
torch.Tensor.backward ([梯度, ...])	计算当前张量相对于图叶子的梯度。
torch.Tensor.detach	返回一个新的 Tensor，与当前图分离。
torch.Tensor.detach_	从创建它的图中分离张量，使其成为叶子节点。
torch.Tensor.register_hook （钩子）	注册反向钩子。
torch.Tensor.register_post_accumulate_grad_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	定义自定义 autograd Function 的前向。
Function.backward	定义使用自动微分反向模式的操作微分公式。
Function.jvp	定义使用前向模式自动微分的不同操作的分步公式。
Function.vmap	定义在 torch.vmap() 下的 autograd.Function 的行为。

上下文方法混合子。

当创建一个新的 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

基础自定义 Function ，用于构建 PyTorch 工具。

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

数值梯度检查 ¶

gradcheck	将通过小有限差分计算得到的梯度与相对于 inputs 中的张量（浮点型或复数类型）的解析梯度进行比较。
gradgradcheck	将通过小有限差分计算得到的梯度与相对于 inputs 和 grad_outputs 中的张量（浮点型或复数类型）的解析梯度进行比较。
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)

管理 autograd 分析器状态和保留结果摘要的上下文管理器。

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

参数:

• enabled（布尔值，可选）- 将此设置为 False 将使此上下文管理器成为空操作。

• use_cuda（布尔值，可选）- 启用使用 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, optional) – 跟踪张量内存分配/释放。

• with_stack (bool, optional) – 记录操作（文件和行号）的源信息。

• with_modules (bool) – 记录操作调用栈对应的模块层次结构（包括函数名）。例如，如果模块 A 的前向调用模块 B 的前向，其中包含 aten::add 操作，则 aten::add 的模块层次结构是 A.B 注意，目前此支持仅适用于 TorchScript 模型，而不适用于 eager 模式模型。

• use_kineto (bool, optional) – 实验，启用 Kineto 分析器的性能分析。

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

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

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

示例

>>> 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 交互式解释器中。

参数:

• enabled (布尔值，可选) – 设置 enabled=False 使此上下文管理器成为空操作。默认： True 。

• record_shapes (布尔值，可选) – 如果 record_shapes=True ，则每个自动微分操作的范围包装将附加有关该操作接收到的 Tensor 参数大小的信息，格式如下： [[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)

前向-后向相关性

在使用 emit_nvtx 在 Nvidia Visual Profiler 中创建的配置文件中查看时，将每个反向传递操作与相应的正向传递操作相关联可能很困难。为了简化这项任务， emit_nvtx 将序列号信息附加到它生成的范围中。

在正向传递过程中，每个函数范围都被 seq=<N> 装饰。 seq 是一个计数器，每次创建一个新的用于反向的 Function 对象并存储起来时，它都会递增。因此，与每个正向函数范围关联的 seq=<N> 注释告诉您，如果由这个正向函数创建反向 Function 对象，则反向对象将接收到序列号 N。在反向传递过程中，每个 C++ 反向 Function 的 apply() 调用的顶层范围都被 stashed seq=<M> 装饰。 M 是反向对象创建时使用的序列号。通过比较反向中的 stashed seq 数字和正向中的 seq 数字，您可以追踪到哪个正向操作创建了每个反向 Function。

在反向传播过程中执行的所有函数也都装饰了 seq=<N> 。在默认的反向传播（使用 create_graph=False ）中，这些信息是无关紧要的，实际上，对于所有这些函数， N 可能简单地是 0。只有与反向函数对象的 apply() 方法关联的顶级范围是有用的，作为将这些函数对象与之前的正向传播相关联的方式。

双向反向传播

相反，如果正在进行带有 create_graph=True 的反向传播（换句话说，如果您正在为双向反向传播做准备），则每个函数在反向传播中的执行都会得到一个非零的有用 seq=<N> 。这些函数本身也可以创建要在双向反向传播中稍后执行的函数对象，就像正向传播中的原始函数一样。反向传播和双向反向传播之间的关系在概念上与正向传播和反向传播之间的关系相同：函数仍然发出带有当前序列号标签的范围，它们创建的函数对象仍然存储这些序列号，在最终的双向反向传播中，函数对象的 apply() 范围仍然带有 stashed seq 标签，这些标签可以与反向传播中的序列号进行比较。

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

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

在使用英特尔(R) VTune Profiler 运行程序时非常有用：

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

仪器和跟踪技术(ITT) API 允许您的应用程序在执行期间生成和控制跨不同英特尔工具的跟踪数据收集。此上下文管理器用于注释英特尔(R) VTune 分析跟踪。借助此上下文管理器，您将能够在英特尔(R) VTune Profiler GUI 中看到标记的范围。

参数:

• 启用（布尔值，可选）- 设置 enabled=False 使此上下文管理器成为空操作。默认： True 。

• record_shapes（布尔值，可选）- 如果 record_shapes=True ，itt 范围包装每个自动微分操作将附加有关该操作接收到的 Tensor 参数大小的信息，格式如下： [[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 跟踪文件并解析 autograd 注释。

调试与异常检测 ¶

class torch.autograd.detect_anomaly(check_nan=True)

上下文管理器，用于为 autograd 引擎启用异常检测。

这做了两件事：

• 启用检测运行前向传递将允许反向传递打印创建失败反向函数的前向操作跟踪。

• 如果 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 上文。

参数:

• 模式（布尔值）- 标志是否启用异常检测（ True ），或禁用（ False ）。

• check_nan（布尔值）- 标志是否在反向生成“nan”时引发错误。

grad_mode.set_multithreading_enabled

设置多线程反向开启或关闭的上下文管理器。

自动微分图 ¶

Autograd 提供了方法，允许用户在反向传播过程中检查图并插入行为。

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

graph.Node.name	返回名称。
graph.Node.metadata	返回元数据。
graph.Node.next_functions
graph.Node.register_hook	注册一个反向钩子。
graph.Node.register_prehook	注册一个反向预钩子。
graph.increment_version	更新 autograd 元数据跟踪给定的 Tensor 是否被就地修改。

一些操作需要在正向传播过程中保存中间结果以便执行反向传播。这些中间结果作为属性保存在 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)[source][source]

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

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

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

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

钩子应具有以下签名：

pack_hook(tensor: Tensor) -> Any

unpack_hook(Any) -> Tensor

pack_hook 的返回值是 unpack_hook 的有效输入。

通常情况下，你希望 unpack_hook(pack_hook(t)) 在值、大小、dtype 和设备方面与 t 相等。

例如：

>>> def pack_hook(x):
...     print("Packing", x)
...     return x
>>>
>>> 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>)

警告

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

警告

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

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

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

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

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

参数:

pin_memory（布尔值）- 如果 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)[source][source]

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

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

参数:

error_message (str) – 当使用已禁用的保存张量默认钩子时，将引发带有此错误信息的 RuntimeError。

返回类型:

生成器[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')[source][source] ¶)

注册多梯度反向钩子。

支持两种模式： "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]
>>>

返回类型:

可移除句柄

class torch.autograd.graph 允许在保存的 tensors 上执行突变[source][source]

允许在反向传播中保存可变张量的上下文管理器。

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

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

返回:

一个存储由该上下文管理器管理的状态的 _AllowMutationOnSavedContext 对象。此对象在调试过程中可能很有用。上下文管理器管理的状态在退出时将自动清除。

返回类型:

生成器[_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)[source][source]

表示 autograd 图中给定梯度边的对象。

要获取计算给定 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)) 。

返回类型:

梯度边缘