# 自动微分包 - torch.autograd > 原文:[`pytorch.org/docs/stable/autograd.html`](https://pytorch.org/docs/stable/autograd.html) `torch.autograd` 提供了实现任意标量值函数自动微分的类和函数。它对现有代码的更改很小 - 您只需要声明应计算梯度的 `Tensor`,并使用 `requires_grad=True` 关键字。目前,我们仅支持浮点 `Tensor` 类型(half、float、double 和 bfloat16)和复数 `Tensor` 类型(cfloat、cdouble)的自动微分。 | `backward` | 计算给定张量相对于图叶子节点的梯度之和。 | | --- | --- | | `grad` | 计算并返回输出相对于输入的梯度之和。 | ## 前向模式自动微分[](#forward-mode-automatic-differentiation "跳转到此标题") 警告 此 API 处于 beta 版本。尽管函数签名很可能不会更改,但在我们将其视为稳定之前,计划增加更多操作符的覆盖范围。 请查看[前向模式 AD 教程](https://pytorch.org/tutorials/intermediate/forward_ad_usage.html)以获取如何使用此 API 的详细步骤。 | `forward_ad.dual_level` | 用于前向 AD 的上下文管理器,在 `dual_level` 上下文中必须进行所有前向 AD 计算。 | | --- | --- | | `forward_ad.make_dual` | 将张量值与其切线关联起来,创建一个用于前向 AD 梯度计算的“双重张量”。 | | `forward_ad.unpack_dual` | 解包“双重张量”,获取其张量值和前向 AD 梯度。| ## 功能高级 API[](#functional-higher-level-api "跳转到此标题") 警告 此 API 处于 beta 版本。尽管函数签名很可能不会更改,但在我们将其视为稳定之前,计划对性能进行重大改进。 本节包含了构建在基本 API 基础上的 autograd 高级 API,允许您计算雅可比矩阵、Hessian 矩阵等。 此 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` | 计算给定函数在输入点处的 Jacobian 矩阵与向量 `v` 的点积。 | | `functional.jvp` | 计算给定函数在输入点处的 Jacobian 矩阵与向量 `v` 的点积。 | | `functional.vhp` | 计算向量`v`与给定标量函数在指定点处的 Hessian 的点积。 | | `functional.hvp` | 计算标量函数的 Hessian 和向量`v`在指定点处的点积。 |## 本地禁用梯度计算[](#locally-disabling-gradient-computation "跳转到此标题的永久链接") 有关无梯度和推断模式之间的区别以及可能与两者混淆的其他相关机制的更多信息,请参阅本地禁用梯度计算。还请参阅本地禁用梯度计算以获取可用于本地禁用梯度的函数列表。## 默认梯度布局 当非稀疏`param`在`torch.autograd.backward()`或`torch.Tensor.backward()`期间接收到非稀疏梯度时,`param.grad`将按以下方式累积。 如果`param.grad`最初为`None`: 1. 如果`param`的内存不重叠且密集,`.grad`将以匹配`param`的步幅创建(从而匹配`param`的布局)。 1. 否则,`.grad` 将以行优先连续步幅创建。 如果`param`已经具有非稀疏的`.grad`属性: 1. 如果`create_graph=False`,`backward()`会原地累积到`.grad`中,从而保留其步幅。 1. 如果`create_graph=True`,`backward()`将用新张量`.grad + new grad`替换`.grad`,尝试(但不保证)匹配现有`.grad`的步幅。 推荐默认行为(在第一次`backward()`之前让`.grad`为`None`,使其布局根据 1 或 2 创建,并根据 3 或 4 保留)以获得最佳性能。调用`model.zero_grad()`或`optimizer.zero_grad()`不会影响`.grad`的布局。 事实上,在每个累积阶段之前将所有的`.grad`重置为`None`,例如: ```py 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`,您的布局也*可能*会被保留。 ## 张量的原地操作[](#in-place-operations-on-tensors "跳转到此标题的永久链接") 在 autograd 中支持原地操作是一个棘手的问题,我们不鼓励在大多数情况下使用它们。Autograd 的积极缓冲区释放和重用使其非常高效,只有在极度内存压力下,才有很少的情况下原地操作实际上会显著降低内存使用量。除非您在极度内存压力下操作,否则您可能永远不需要使用它们。 ### 原地正确性检查[](#in-place-correctness-checks "跳转到此标题的永久链接") 所有`Tensor`都会跟踪应用于它们的原地操作,如果实现检测到一个张量在其中一个函数中保存以进行反向传播,但之后在原地进行修改,那么一旦开始反向传播,就会引发错误。这确保了如果您使用原地函数而没有看到任何错误,您可以确信计算的梯度是正确的。 ## 变量(已弃用) 警告 Variable API 已被弃用:不再需要使用 Variables 来使用带有 `requires_grad` 设置为 `True` 的张量进行自动求导。自动求导自动支持 `requires_grad` 设置为 `True` 的张量。以下是有关已更改内容的快速指南: + `Variable(tensor)` 和 `Variable(tensor, requires_grad)` 仍然按预期工作,但它们返回张量而不是 Variables。 + `var.data` 和 `tensor.data` 是相同的。 + 现在,`var.backward()、var.detach()、var.register_hook()` 等方法在具有相同方法名称的张量上也可以使用。 此外,现在可以使用工厂方法(如 `torch.randn()`、`torch.zeros()`、`torch.ones()` 等)创建 `requires_grad=True` 的张量,如下所示: `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`。 | ## Function ```py class torch.autograd.Function(*args, **kwargs)¶ ``` 用于创建自定义自动求导函数的基类。 要创建自定义自动求导函数,请继承此类并实现 `forward()` 和 `backward()` 静态方法。然后,在前向传播中使用您的自定义操作,调用类方法 `apply`。不要直接调用 `forward()`。 为了确保正确性和最佳性能,请确保在 `ctx` 上调用正确的方法,并使用 `torch.autograd.gradcheck()` 验证您的反向函数。 有关如何使用此类的更多详细信息,请参阅 扩展 torch.autograd。 示例: ```py >>> 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()`下的行为。 | ## 上下文方法混合 创建新的`Function`时,以下方法可用于 ctx。 | `function.FunctionCtx.mark_dirty` | 将给定张量标记为原位操作中修改的张量。 | | --- | --- | | `function.FunctionCtx.mark_non_differentiable` | 将输出标记为不可微分。 | | `function.FunctionCtx.save_for_backward` | 保存给定张量,以便将来调用`backward()`。 | | `function.FunctionCtx.set_materialize_grads` | 设置是否实现梯度张量。 | ## 数值梯度检查[](#module-torch.autograd.gradcheck "Permalink to this heading") | `gradcheck` | 检查通过小的有限差分计算的梯度与`inputs`中的浮点或复数类型张量相对于解析梯度的梯度,这些张量具有`requires_grad=True`。 | | --- | --- | | `gradgradcheck` | 检查通过小的有限差分计算的梯度的梯度与`inputs`和`grad_outputs`中的浮点或复数类型张量相对于解析梯度的梯度,这些张量具有`requires_grad=True`。 | ## 性能分析器 Autograd 包括一个性能分析器,让您可以检查模型内不同运算符的成本 - 包括在 CPU 和 GPU 上。目前实现了三种模式 - 仅 CPU 使用`profile`。基于 nvprof 的(同时注册 CPU 和 GPU 活动)使用`emit_nvtx`。基于 vtune 性能分析器的使用`emit_itt`。 ```py 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, use_mtia=False, experimental_config=None)¶ ``` 上下文管理器,管理自动求导分析器状态并保存结果摘要。 在底层,它只记录在 C++中执行的函数事件,并将这些事件暴露给 Python。您可以将任何代码包装在其中,它只会报告 PyTorch 函数的运行时。注意:性能分析器是线程本地的,并自动传播到异步任务中 参数 + **enabled**(*bool*,可选)- 将其设置为 False 会使此上下文管理器无效。 + **use_cuda**(*bool*,可选)- 启用使用 cudaEvent API 对 CUDA 事件进行计时。每个张量操作大约增加 4 微秒的开销。 + **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 模型,而不适用于急切模式模型。 + **use_kineto**(*bool*) – 实验性的,启用使用 Kineto 分析器进行分析。 + **use_cpu**(*bool*) – 分析 CPU 事件;将其设置为`False`需要`use_kineto=True`,可用于降低仅针对 GPU 的分析的开销。 + **experimental_config**(*_ExperimentalConfig*) – 由像 Kineto 这样的分析器库使用的一组实验选项。请注意,不保证向后兼容性。 示例 ```py >>> 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` | 将事件列表导出为 Chrome 跟踪工具文件。 | | --- | --- | | `profiler.profile.key_averages` | 对所有函数事件按其键进行平均。 | | `profiler.profile.self_cpu_time_total` | 返回在 CPU 上花费的总时间。 | | `profiler.profile.total_average` | 对所有事件进行平均。 | ```py class torch.autograd.profiler.emit_nvtx(enabled=True, record_shapes=False)¶ ``` 使每个自动求导操作发出一个 NVTX 范围的上下文管理器。 在使用 nvprof 运行程序时很有用: ```py nvprof --profile-from-start off -o trace_name.prof -- ``` 不幸的是,没有办法强制 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 范围将附加有关该操作接收的 Tensor 参数大小的信息,格式如下:`[[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...]` 非张量参数将用`[]`表示。参数将按照后端操作接收它们的顺序列出。请注意,此顺序可能与这些参数在 Python 端传递的顺序不匹配。还请注意,形状记录可能会增加 nvtx 范围创建的开销。默认值:`False` 示例 ```py >>> 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`会向其生成的范围附加序列号信息。 在前向传递期间,每个函数范围都被装饰为`seq=`。`seq`是一个运行计数器,每次创建一个新的反向 Function 对象并存储用于反向传递时递增。因此,与每个前向函数范围相关联的`seq=`注释告诉您,如果此前向函数创建了一个反向 Function 对象,那么反向对象将接收序列号 N。在反向传递期间,包装每个 C++反向 Function 的`apply()`调用的顶级范围都被装饰为`stashed seq=`。`M`是创建反向对象时的序列号。通过在反向传递中比较反向中的`stashed seq`号和前向中的`seq`号,您可以追踪哪个前向操作创建了每个反向 Function。 在反向传递期间执行的任何函数也都被装饰为`seq=`。在默认反向传递(`create_graph=False`)中,此信息是无关紧要的,实际上,对于所有这样的函数,`N`可能只是 0。只有与反向 Function 对象的`apply()`方法相关联的顶级范围是有用的,作为将这些 Function 对象与之前的前向传递相关联的一种方式。 **双向传递** 另一方面,如果正在进行带有`create_graph=True`的反向传递(换句话说,如果您正在为双向传递做准备),则在反向期间执行的每个函数都会获得一个非零、有用的`seq=`。这些函数本身可能会创建 Function 对象,以便稍后在双向传递期间执行,就像前向传递中的原始函数一样。反向和双向之间的关系在概念上与前向和反向之间的关系相同:函数仍然发出当前序列号标记的范围,它们创建的 Function 对象仍然存储这些序列号,并且在最终的双向传递期间,Function 对象的`apply()`范围仍然标记有`stashed seq`号,可以与反向传递中的 seq 号进行比较。 ```py class torch.autograd.profiler.emit_itt(enabled=True, record_shapes=False)¶ ``` 上下文管理器,使每个自动求导操作都发出一个 ITT 范围。 在运行程序时在英特尔(R) VTune Profiler 下很有用: ```py vtune <--vtune-flags> ``` 仪器和跟踪技术(ITT)API 使您的应用程序能够在其执行过程中生成和控制跨不同英特尔工具的跟踪数据的收集。这个上下文管理器是用来注释英特尔(R) VTune Profiling 跟踪的。借助这个上下文管理器,您将能够在英特尔(R) VTune Profiler GUI 中看到标记的范围。 参数 + **enabled** ([*bool*](https://docs.python.org/3/library/functions.html#bool "(在 Python v3.12 中)")*,* *可选*) – 设置 `enabled=False` 会使这个上下文管理器无效。默认值:`True`。 + **record_shapes** ([*bool*](https://docs.python.org/3/library/functions.html#bool "(在 Python v3.12 中)")*,* *可选*) – 如果 `record_shapes=True`,则每个自动求导操作周围的 itt 范围将附加关于该操作接收到的张量参数大小的信息,格式如下:`[[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...]` 非张量参数将用 `[]` 表示。参数将按照它们被后端操作接收的顺序列出。请注意,这个顺序可能与这些参数在 Python 端传递的顺序不匹配。还请注意,形状记录可能会增加 itt 范围创建的开销。默认值:`False` 示例 ```py >>> with torch.autograd.profiler.emit_itt(): ... model(x) ``` | `profiler.load_nvprof` | 打开一个 nvprof 跟踪文件并解析自动求导注释。 | | --- | --- | ## 异常检测 ```py class torch.autograd.detect_anomaly(check_nan=True)¶ ``` 上下文管理器,用于启用自动求导引擎的异常检测。 这样做两件事: + 在启用检测的情况下运行前向传递将允许反向传递打印创建失败的反向函数的前向操作的回溯。 + 如果 `check_nan` 为 `True`,任何生成“nan”值的反向计算都将引发错误。默认为 `True`。 警告 此模式应仅用于调试,因为不同的测试会减慢程序执行速度。 示例 ```py >>> 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 "", line 1, in 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 "", 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 out = run_fn(inp) File "tmp.py", line 44, in run_fn out = MyFunc.apply(a) Traceback (most recent call last): File "", line 4, in 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 "", line 8, in backward RuntimeError: Some error in backward ``` ```py class torch.autograd.set_detect_anomaly(mode, check_nan=True)¶ ``` 上下文管理器,用于设置自动求导引擎的异常检测开启或关闭。 `set_detect_anomaly` 将根据其参数 `mode` 启用或禁用自动求导异常检测。它可以作为上下文管理器或函数使用。 有关异常检测行为的详细信息,请参阅上面的 `detect_anomaly`。 参数 + **mode** ([*bool*](https://docs.python.org/3/library/functions.html#bool "(在 Python v3.12 中)")*) – 标志,指示是否启用异常检测 (`True`) 或禁用 (`False`)。 + **check_nan** ([*bool*](https://docs.python.org/3/library/functions.html#bool "(在 Python v3.12 中)")) – 标志,指示在反向传播生成“nan”时是否引发错误。 ## 自动求导图 自动求导公开了允许检查图并在反向传播期间插入行为的方法。 如果张量是由自动求导记录的操作的输出(即启用了 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` | 注册一个反向钩子。 | | `graph.Node.register_prehook` | 注册一个反向预钩子。 | 一些操作需要在前向传递期间保存中间结果以执行反向传递。这些中间结果保存为`grad_fn`上的属性,并可以访问。例如: ```py >>> 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,但在训练期间不适合,则这是非常有用的。另请参阅保存张量的钩子。 ```py class torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook)¶ ``` 为保存的张量设置一对 pack / unpack 钩子的上下文管理器。 使用这个上下文管理器来定义操作的中间结果在保存前应如何打包,并在检索时解包。 在这种情况下,每次操作保存张量进行反向传递时(包括使用`save_for_backward()`保存的中间结果以及由 PyTorch 定义的操作记录的结果),都会调用`pack_hook`函数。然后,`pack_hook`的输出将存储在计算图中,而不是原始张量。 当需要访问保存的张量时,即在执行`torch.Tensor.backward()`或`torch.autograd.grad()`时,将调用`unpack_hook`。它以`pack_hook`返回的*packed*对象作为参数,并应返回一个与原始张量内容相同的张量(作为输入传递给相应的`pack_hook`)。 钩子应具有以下签名: > pack_hook(tensor: Tensor) -> Any > > unpack_hook(Any) -> Tensor 其中`pack_hook`的返回值是`unpack_hook`的有效输入。 一般来说,您希望`unpack_hook(pack_hook(t))`在值、大小、dtype 和设备方面等于`t`。 示例: ```py >>> 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=) >>> y.sum().backward() Unpacking tensor([1., 1., 1., 1., 1.], requires_grad=True) Unpacking tensor([2., 2., 2., 2., 2.], grad_fn=) ``` 警告 对输入执行就地操作到任一钩子可能导致未定义行为。 警告 一次只允许一个钩子对。当递归嵌套这个上下文管理器时,只有最内层的一对钩子将被应用。 ```py class torch.autograd.graph.save_on_cpu(pin_memory=False, device_type='cuda')¶ ``` 在前向传递期间保存的张量将存储在 CPU 上,然后在反向传递时检索。 在这个上下文管理器中执行操作时,在前向传递期间保存在图中的中间结果将被移动到 CPU,然后在需要进行反向传递时复制回原始设备。如果图已经在 CPU 上,则不执行张量复制。 使用这个上下文管理器来在计算和 GPU 内存使用之间进行交换(例如,在训练期间模型不适合 GPU 内存时)。 参数 **pin_memory**([*bool*](https://docs.python.org/3/library/functions.html#bool "(in Python v3.12)")) - 如果为`True`,张量将在打包期间保存到 CPU 固定内存,并在解包期间异步复制到 GPU。默认为`False`。另请参阅使用固定内存缓冲区。 示例: ```py >>> 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 ``` ```py class torch.autograd.graph.disable_saved_tensors_hooks(error_message)¶ ``` 禁用保存的张量默认钩子功能的上下文管理器。 如果您正在创建一个与保存的张量默认钩子不兼容的特性,则此功能很有用。 参数 **error_message**([*str*](https://docs.python.org/3/library/stdtypes.html#str "(in Python v3.12)")) - 当使用保存的张量默认钩子时,当它们被禁用时,会引发带有此错误消息的 RuntimeError。 示例: ```py >>> 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 ``` ```py class torch.autograd.graph.register_multi_grad_hook(tensors, fn)¶ ``` 注册一个多梯度反向钩子。 在计算完与`tensors`中的每个张量相关的梯度后,将调用钩子。如果一个张量在`tensors`中但不是图的一部分,或者一个张量不需要计算当前`.backward()`或`.grad()`调用中指定的任何`inputs`的梯度,那么这个张量将被忽略,钩子将不会等待其梯度被计算。 在计算完每个未被忽略的张量的梯度后,将调用`fn`并传递这些梯度。对于那些没有计算梯度的张量,将传递`None`。 钩子不应修改其参数。 此函数返回一个具有`handle.remove()`方法的句柄,用于移除钩子。 注意 有关此钩子的执行时间以及其执行顺序与其他钩子的关系的更多信息,请参见反向钩子执行。 示例: ```py >>> 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] >>> ``` ```py class torch.autograd.graph.allow_mutation_on_saved_tensors¶ ``` 允许变异保存用于反向传播的张量的上下文管理器。 在这个上下文管理器下,保存用于反向传播的张量在变异时会被克隆,因此原始版本仍然可以在反向传播期间使用。通常,对保存用于反向传播的张量进行变异会导致在反向传播期间使用时引发错误。 为了确保正确的行为,前向和后向都应该在相同的上下文管理器下运行。 返回 一个存储由这个上下文管理器管理的状态的 _AllowMutationOnSavedContext 对象。这个对象对于调试目的很有用。上下文管理器管理的状态在退出时会自动清除。 示例: ```py >>> 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=) ``` ```py class torch.autograd.graph.GradientEdge(node, output_nr)¶ ``` 表示 autograd 图中给定梯度边缘的对象。要获取将计算给定张量梯度的梯度边缘,可以执行`edge = autograd.graph.get_gradient_edge(tensor)`。 ```py torch.autograd.graph.get_gradient_edge(tensor)¶ ``` 获取给定张量的梯度的梯度边缘。 特别地,调用`g = autograd.grad(loss, input)`和`g = autograd.grad(loss, get_gradient_edge(input))`是等价的。