目录

自动微分包 - torch.autograd

torch.autograd 提供了实现任意标量值函数自动微分的类和函数。它需要对现有代码进行最小的更改 - 你只需要为需要计算梯度的 Tensor 使用 requires_grad=True 关键字声明。到目前为止,我们只支持浮点 Tensor 类型(half, float, double 和 bfloat16)和复数 Tensor 类型(cfloat, cdouble)的自动微分。

backward

计算给定张量相对于图叶子的梯度之和。

grad

计算并返回输出相对于输入的梯度之和。

前向模式自动微分

警告

此 API 处于测试阶段。虽然函数签名极不可能更改,但在将其视为稳定之前,计划增加操作覆盖率。

请参阅前向模式自动微分教程 以获取有关如何使用此 API 的详细步骤。

forward_ad.dual_level

启用前向自动微分的上下文管理器。

forward_ad.make_dual

将一个张量值与前向梯度(即切线)相关联,以创建“双张量”,用于计算前向自动微分(AD)梯度。

forward_ad.unpack_dual

将一个“双重张量”解包,以获取其张量值及其前向 AD 梯度。

高级功能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 与给定标量函数在输入点处的海森矩阵之间点积的函数。

functional.hvp

计算给定标量函数的Hessian矩阵与由输入给出的点处的向量 v 的点积的函数。

局部禁用梯度计算

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

默认梯度布局

当一个非稀疏的 param 在进行 torch.autograd.backward()torch.Tensor.backward() param.grad 时,梯度累积如下。

如果 param.grad 初始为 None

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

  2. 否则,.grad 会以行主序连续 strides 创建。

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

  1. 如果 create_graph=Falsebackward() 累积为 .grad 原地进行,这将保留其步长。

  2. 如果用create_graph=Truebackward()替换.grad,并生成一个新的张量.grad + new grad,它会尝试(但不保证)匹配现有的.grad的步长。

默认行为(在首次出现之前让 .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。 只要 create_graph=False 存在,3 就能保证你的布局得以保留。 4 表示即使 create_graph=True,你的布局也很可能被保留。

就地操作在张量上

支持就地操作在自动微分中是一个复杂的问题,我们不建议在大多数情况下使用它们。自动微分的激进缓冲区释放和重用使其非常高效,实际上能够显著降低内存使用的场合很少见。除非你在极高的内存压力下运行,否则你可能永远不需要使用它们。

原地正确性检查

所有 Tensor 都会跟踪对其执行的就地操作, 如果实现检测到在其中一个函数中保存了张量以备反向传播,但在之后进行了就地修改,则一旦开始反向传播就会引发错误。 这确保了如果你使用就地函数且没有看到任何错误,可以确定计算出的梯度是正确的。

变量(已弃用)

警告

Variable API已被弃用:变量不再需要用于张量的自动微分。Autograd 自动支持设置为 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_gradFalse的张量将按照惯例被视为叶张量。

torch.Tensor.backward([gradient, ...])

计算当前张量的梯度。

torch.Tensor.detach

返回一个新的张量,与当前计算图分离。

torch.Tensor.detach_

分离该张量与其创建图的关系,使其成为一个叶节点。

torch.Tensor.register_hook(钩子)

注册一个反向传播钩子。

torch.Tensor.retain_grad()

启用此张量在其 grad 中填充数据 backward() 期间。

功能

class torch.autograd.Function(*args, **kwargs)[source]

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

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

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

参见扩展 torch.autograd以了解更多关于如何使用此类的详细信息。

Examples:

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

定义一个公式,用于对操作进行反向模式自动微分的求导(vjp 函数的别名)。

Function.jvp

定义了一个用于对前向模式自动微分操作进行求导的公式。

上下文方法混入

当创建一个新的 Function 时,以下方法可用于ctx

function.FunctionCtx.mark_dirty

将给定的张量标记为在原地操作中被修改。

function.FunctionCtx.mark_non_differentiable

将输出标记为不可微分。

function.FunctionCtx.save_for_backward

保存给定的张量,以供将来调用 backward()

function.FunctionCtx.set_materialize_grads

设置是否生成输出梯度张量。

数值梯度检查

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, record_shapes=False, with_flops=False, profile_memory=False, with_stack=False, with_modules=False, use_kineto=False, use_cpu=True, experimental_config=None)[source]

用于管理自动求导分析器状态并保存结果摘要的上下文管理器。 在底层,它只是记录在 C++ 中执行函数的事件,并将这些事件暴露给 Python。你可以将其封装任何代码,并且它只会报告 PyTorch 函数的运行时间。 注意:分析器是线程本地的,并会自动传播到异步任务中

Parameters:
  • 启用 (布尔值, 可选) – 将此设置为 False 会使此上下文管理器无效。

  • use_cuda (bool, optional) – 启用使用 cudaEvent API 对 CUDA 事件进行计时。每个张量操作会增加大约 4 微秒的开销。

  • record_shapes (bool, 可选) – 如果启用了形状记录,将收集有关输入维度的信息。这允许您查看哪些维度在后台被使用,并进一步通过使用prof.key_averages(group_by_input_shape=True)按它们进行分组。请注意,形状记录可能会使您的性能分析数据产生偏差。建议分别运行带有和不带形状记录的程序以验证时间。对于嵌套函数调用中最底层的事件,偏差可能可以忽略不计。但对于更高层次的函数,总自 CPU 时间可能会因形状收集而人为增加。

  • with_flops (bool, 可选) – 如果设置了 with_flops,分析器将使用操作符的输入形状估算 浮点运算次数(FLOPs)。这使得可以估算硬件性能。目前, 此选项仅适用于矩阵乘法和二维卷积操作。

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

  • with_stack (bool, 可选) – 记录操作的来源信息(文件和行号)。

  • with_modules (bool) – 记录模块层次结构(包括函数名称) 对应于操作调用栈。例如,如果模块 A 的前向调用中包含模块 B 的前向调用, 而模块 B 中有一个 aten::add 操作, 那么 aten::add 的模块层次结构为 A.B。 请注意,目前此功能仅支持 TorchScript 模型,而不支持即时模式模型。

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

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

  • 实验性配置 (_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

对所有事件求平均值。

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

上下文管理器,使每个自动微分操作发出一个 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中。

Parameters:
  • 启用 (布尔值, 可选) – 设置 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)

前向后向相关性

当使用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。只有与反向传播 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)[source]

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

在英特尔(Intel)VTune分析器下运行程序时很有用:

vtune <--vtune_flags> <regular command here>

仪器仪表和追踪技术(ITT)API使您的应用程序能够在执行过程中生成和控制不同英特尔工具之间的跟踪数据收集。 此上下文管理器用于标注英特尔® VTune 配置文件跟踪。借助此上下文管理器的帮助,您将在英特尔® VTune 配置文件 GUI 中看到带标签的范围。

Parameters:
  • 启用 (布尔值, 可选) – 设置 enabled=False 会使此上下文管理器无效。 默认值: True

  • record_shapes (bool, 可选) – 如果为 record_shapes=True,每个autograd操作的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 跟踪文件并解析 autograd 注释。

异常检测

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

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

这做了两件事:

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

  • 如果 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 以了解异常检测行为的详细信息。

Parameters:
  • 模式 (布尔值) – 标记是否启用异常检测 (True), 或禁用 (False)。

  • 检查_nan (bool) – 标记是否在反向传播生成“nan”时触发错误

已保存的张量默认钩子

某些操作需要在前向传递过程中保存中间结果,以便执行后向传递。 您可以使用挂钩定义这些保存的张量应该如何打包/解包。 一个常见的应用是通过将这些中间结果保存到磁盘或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()时。它以packed对象作为参数,该对象由pack_hook返回,并应返回一个与原始张量内容相同的张量(作为输入传递给相应的pack_hook)。

钩子应该具有以下签名:

pack_hook(tensor: Tensor) -> Any

unpack_hook(Any) -> Tensor

其中返回值 pack_hookunpack_hook 的有效输入。

一般来说,您希望unpack_hook(pack_hook(t))在值、大小、数据类型和设备方面等于t

Example:

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

上下文管理器,其中由前向传播保存的张量将被存储在 CPU 上,然后用于反向传播。

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

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

Parameters:

pin_memory (布尔值) – 如果设置为 True,张量将在打包时保存到CPU的固定内存中,并在解包时异步复制到GPU。 默认值为 False。 另请参阅 使用固定的内存缓冲区

Example:

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

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

如果你正在创建一个不支持保存张量默认钩子的功能,这可能会对你有用。

Parameters:

错误消息 (字符串) – 当保存的张量默认钩子被使用而它们已被禁用时,会引发带有此错误消息的 RuntimeError。

Example:

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

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

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

查看教程

资源

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

查看资源