• 文档 >
  • 自动微分包 - torch.autograd
快捷方式

自动微分包 - torch.autograd

torch.autograd 提供了实现任意标量值函数自动微分的类和函数。它对现有代码的要求极低 - 你只需使用 Tensor 关键字 requires_grad=True 声明需要计算梯度的张量即可。目前,我们仅支持浮点 Tensor 类型(half、float、double 和 bfloat16)和复数 Tensor 类型(cfloat、cdouble)的自动微分。

backward

计算给定张量关于图叶的梯度和。

grad

计算并返回输出关于输入的梯度和。

前向模式自动微分

警告

此 API 处于测试阶段。即使函数签名极不可能更改,我们仍计划在将此 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() 返回的命名元组,包含双张量的原始和切线分量。

功能性高级 API

警告

此 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

计算给定标量函数的海森矩阵。

functional.vjp

计算向量 v 与给定函数在输入给出的点处的雅可比行列式之间的点积。

functional.jvp

计算给定函数在输入给出的点处的雅可比行列式与向量 v 之间的点积。

functional.vhp

计算向量 v 与给定标量函数在指定点处的海森矩阵之间的点积。

functional.hvp

计算标量函数的海森矩阵与向量 v 在指定点处的点积。

局部禁用梯度计算

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

默认梯度布局

当非稀疏 paramtorch.autograd.backward()torch.Tensor.backward() 期间接收到非稀疏梯度时,param.grad 会累积如下。

如果 param.grad 最初为 None

  1. 如果 param 的内存非重叠且密集,则会使用与 param 匹配的步长创建 .grad(因此匹配 param 的布局)。

  2. 否则,将使用行主连续步长创建 .grad

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

  1. 如果 create_graph=False,则 backward() 会就地累积到 .grad 中,从而保留其步长。

  2. 如果 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,例如

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 已弃用:不再需要变量才能将 autograd 与张量一起使用。Autograd 自动支持将 requires_grad 设置为 True 的张量。以下请找到有关已更改内容的快速指南

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

  • var.datatensor.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)

张量 autograd 函数

torch.Tensor.grad

此属性默认情况下为 None,并且在首次调用 backward()self 计算梯度时变为张量。

torch.Tensor.requires_grad

如果需要为该张量计算梯度,则为 True,否则为 False

torch.Tensor.is_leaf

根据惯例,所有 requires_gradFalse 的张量都将是叶张量。

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

使此张量能够在 grad 中填充 backward()

函数

class torch.autograd.Function(*args, **kwargs)[源代码]

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

要创建自定义 autograd.Function,请继承此类并实现 forward()backward() 静态方法。然后,要在前向传播中使用自定义 op,请调用类方法 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 的 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

设置是否实现梯度张量。

自定义函数实用程序

backward 方法的装饰器。

function.once_differentiable

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

function.BackwardCFunction

此类用于内部自动微分工作。

function.InplaceFunction

此类仅出于向后兼容性原因而存在。

function.NestedIOFunction

此类仅出于向后兼容性原因而存在。

数值梯度检查

gradcheck

检查通过小有限差分计算的梯度与 inputs 中的浮点或复数类型张量相对于分析梯度的梯度,并且 requires_grad=True

gradgradcheck

检查通过小有限差分计算的梯度的梯度与 inputsgrad_outputs 中的浮点或复数类型张量相对于分析梯度的梯度,并且 requires_grad=True

GradcheckError

gradcheck()gradgradcheck() 引发的错误。

分析器

Autograd 包含一个分析器,它允许您检查模型中不同运算符的成本 - 同时在 CPU 和 GPU 上。目前实现了三种模式 - 仅使用 CPU profile。基于 nvprof(同时记录 CPU 和 GPU 活动)使用 emit_nvtx。以及基于 vtune 分析器使用 emit_itt

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, 可选) – 启用 CUDA 事件计时,同时使用 cudaEvent API。为每个张量操作增加大约 4us 的开销。

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

  • with_flops (bool, 可选) – 如果设置了 with_flops,分析器将使用运算符的输入形状估算 FLOP(浮点运算)值。这允许人们估算硬件性能。目前,此选项仅适用于矩阵乘法和 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 等分析器库使用的一组实验选项。请注意,不保证向后兼容性。

示例

>>> 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)[source]

上下文管理器,使每个 autograd 操作都发出 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,则包装每个 autograd op 的 nvtx 范围将附加该 op 接收的张量参数大小的信息,格式如下:[[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...] 非张量参数将由 [] 表示。参数将按后端 op 接收它们的顺序列出。请注意,此顺序可能与在 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 创建的概要文件时,将每个反向传递 op 与相应的正向传递 op 关联起来可能很困难。为了简化此任务,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>。这些函数本身可能会创建函数对象,以便在双反向期间稍后执行,就像正向传递中的原始函数一样。反向和双反向之间的关系在概念上与正向和反向之间的关系相同:函数仍然会发出当前序列号标记的范围,它们创建的函数对象仍然会存储这些序列号,并且在最终的双反向期间,函数对象的 apply() 范围仍然标记有 stashed seq 编号,可以将其与反向传递中的 seq 编号进行比较。

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

上下文管理器,使每个自动梯度操作发出一个 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, 可选) – 设置 enabled=False 使此上下文管理器变为无操作。默认值:True

  • record_shapes (bool, optional) – 如果 record_shapes=True,则包装每个 autograd op 的 itt 范围将追加该 op 收到的张量参数大小的信息,格式如下: [[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...] 非张量参数将由 [] 表示。参数将按后端 op 收到的顺序列出。请注意,此顺序可能与在 Python 端传递这些参数的顺序不匹配。另请注意,形状记录可能会增加 itt 范围创建的开销。默认值: False

示例

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

profiler.load_nvprof

打开 nvprof 跟踪文件并解析 autograd 注释。

调试和异常检测

class torch.autograd.detect_anomaly(check_nan=True)[source]

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

这会做两件事

  • 在启用检测的情况下运行前向传递将允许反向传递打印创建失败反向函数的前向操作的回溯。

  • 如果 check_nanTrue,则生成“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)[source]

上下文管理器,用于打开或关闭自动梯度引擎的异常检测。

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

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

参数
  • mode (bool) – 标志是否启用异常检测 (True) 或禁用 (False)。

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

grad_mode.set_multithreading_enabled

上下文管理器,用于打开或关闭多线程反向传播。

自动梯度图

自动梯度公开了一些方法,允许人们检查图并在反向传播期间插入行为。

如果张量是记录在自动梯度中的操作的输出(即,启用了 grad_mode 并且至少一个输入需要梯度),则 torch.Tensorgrad_fn 属性会保存一个 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)[source]

上下文管理器,为已保存的张量设置一对打包/解包挂钩。

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

在此上下文中,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
>>>
>>> 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]

上下文管理器,在该管理器下,前向传播保存的张量将存储在 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)[source]

禁用已保存张量默认挂钩特性的上下文管理器。

如果您正在创建不适用于已保存张量默认挂钩的功能,则很有用。

参数

error_message (str) – 当已禁用已保存张量默认挂钩时使用它们时,将引发带有此错误消息的 RuntimeError。

示例

>>> 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
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]
>>>
torch.autograd.graph.allow_mutation_on_saved_tensors[源代码]

在允许对为反向保存的张量进行变异的上下文管理器下。

在此上下文管理器下,为反向传播保存的张量在突变时被克隆,因此原始版本仍可在反向传播期间使用。通常,突变为反向传播保存的张量将在反向传播期间使用时导致错误。

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

返回

存储此上下文管理器管理的状态的 _AllowMutationOnSavedContext 对象。此对象可用于调试目的。上下文管理器管理的状态在退出时会自动清除。

示例

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

文档

获取 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源