自动微分包 - torch.autograd

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

它只需要对现有代码进行少量修改——只需使用requires_grad=True关键字为需要计算梯度的Tensor声明即可。目前,我们仅支持浮点型Tensor类型(half, float, double 和 bfloat16)和复数型Tensor类型的自动微分。

backward

计算给定张量相对于图中叶子节点的梯度总和。

grad

计算并返回输出对输入的梯度总和。

前向模式自动 differentiation

或者更自然表达:

前向自动微分

由于“differentiation”在数学中通常翻译为“微分”,这里采用“前向自动微分”更为简洁和常用。如果需要保持与英文原文的直接对应,可以使用第一种方式。

警告

此 API 处于测试阶段。虽然函数签名几乎不会改变,但我们计划在正式认定其为稳定版本之前,进一步提高操作符的覆盖率。

请参阅前向模式AD教程,以了解如何使用此API的详细步骤。

forward_ad.dual_level

这是一个用于正向自动微分的上下文管理器,在此过程中所有的正向自动微分计算都必须在dual_level上下文中完成。

forward_ad.make_dual

将张量值与其切线关联,形成“双重张量”,以进行正向自动微分的梯度计算。

forward_ad.unpack_dual

解开一个 "双张量",以便同时获得其 Tensor 值和正向自动微分的梯度。

forward_ad.enter_dual_level

输入一个新的正向传播梯度值。

torch.autograd.forward_ad.exit_dual_level

退出正向梯度级别。

torch.autograd.forward_ad.UnpackedDualTensor

unpack_dual() 返回的命名元组,包含了对偶张量的原语分量和切线分量。

高级功能API

警告

此 API 处于测试阶段。虽然函数签名几乎不会改变,但我们计划在正式认定其为稳定版本之前进行重大性能改进。

本节包含了基于上述基本API的autograd高级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在指定点的点积。

本地禁用梯度计算

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

默认渐变布局

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

如果 param.grad 初始为 None:

  1. 如果param的内存是非重叠且密集的,那么.grad会以与param相同的步长(即与param的布局一致)进行创建。

  2. 否则,.grad 将采用行优先的连续步长进行创建。

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

  1. 如果 create_graph=False,则 backward() 会就地累加到 .grad,以保持其步长不变。

  2. 如果 create_graph=Truebackward() 会将 .grad 替换为一个新的张量 .grad + new grad,尝试(但不能保证)匹配现有 .grad 的 strides。

默认情况下,在第一次调用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 的 strides,在第一次调用 backward() 之前,将 param.grad = 赋值为具有所需 strides 的零张量,并且永远不要将其重置为 None。只要 create_graph=False,3 就能保证你的布局会被保留。4 表示即使 create_graph=True,你的布局也很可能会被保留。

张量的就地操作

在自动微分中支持就地操作是一个复杂的问题,我们不建议在大多数情况下使用这些操作。Autograd 的激进缓冲区释放和重用使其非常高效,通常情况下,就地操作并不会显著降低内存使用量。除非你处于严重的内存压力下,否则你可能根本不需要使用它们。

就地验证检查

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

变量( deprecated )

警告

Variable API 已被弃用:使用 autograd 时不再需要 Variables。Autograd 自动支持将 requires_grad 设置为 True 的 Tensor。以下是一份变更的快速指南:

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

张量自动求导函数

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

允许此张量在调用backward()时,grad属性被填充。

函数

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

torch.vmap() 中定义此 autograd.Function 的行为。

context 方法混合

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

function.FunctionCtx.mark_dirty

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

FunctionCtx.mark_non_differentiable

将输出标记为不可微分。

function.FunctionCtx.save_for_backward

为将来调用 backward() 保存给定的张量。

torch.autograd.function.FunctionCtx.set_materialize_grads

设置是否将梯度张量化。

自定义功能工具

用于反向方法的装饰器。

torch.autograd.function.once_differentiable

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

BackwardCFunction

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

InplaceFunction

此类仅为了向后兼容而存在。

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,以及使用 emit_itt 的 vtune 剖析器。

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)[源代码]

管理自动梯度记录器状态的上下文管理器,并保存结果摘要。

实际上,它只是记录了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(浮点运算)值。这有助于评估硬件性能。目前,此选项仅适用于矩阵乘法和二维卷积操作。

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

  • with_stack (bool, optional) – 是否记录操作的源代码位置(文件名和行号)。

  • with_modules (bool) – 记录操作调用栈对应的模块层次结构(包括函数名称)。例如,如果模块 A 的前向方法调用了模块 B 的前向方法,并且该方法包含一个 aten::add 操作,则 aten::add 的模块层次结构为 A.B。需要注意的是,此功能目前仅支持 TorchScript 模型,而不支持 eager mode 模型。

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

  • use_cpu (bool, optional) – 是否采集 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 tracing 工具支持的文件格式。

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

torch.autograd.profiler_util.Kernel

torch.autograd.profiler_util.MemRecordsAcc

用于在区间内访问 mem_records 的加速结构。

torch.autograd.profiler_util.StringTable

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 视觉调试器(nvvp)来可视化时间线,或者通过调用 torch.autograd.profiler.load_nvprof() 在 Python REPL 中加载和检查结果。

参数
  • enabled (bool, 可选) – 当设置为 enabled=False 时,此上下文管理器将不执行任何操作。默认值:True

  • record_shapes (bool, 可选) – 如果设置为 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 是一个递增计数器,在每次创建新的反向 Function 对象时增加,并存储起来用于后续的反向传播。因此,与每个正向函数范围关联的seq=<N>注释表示如果该正向函数创建了一个新的反向 Function 对象,则该对象将被赋予序列号 N。在反向传播过程中,封装每个 C++ 反向 Function 的 apply() 调用的顶级范围会被标记为stashed seq=<M>M 是反向对象创建时所使用的序列号。通过比较反向传播中的stashed seq 数字与正向传播中的seq数字,你可以追踪到每个反向 Function 对象是由哪个正向操作创建的。

在反向传递过程中执行的函数会被标记为seq=<N>。默认情况下(即create_graph=False),这些信息是无关紧要的,对于所有这样的函数来说,N 可能都是 0。只有与 backward Function 对象的 apply() 方法相关的顶级范围才有用,以便将这些 Function 对象与其之前的前向传递关联起来。

双重反向传播

另一方面,如果正在进行带有create_graph=True参数的反向传播(换句话说,如果你正在为二次反向传播做准备),那么在反向传播期间每个函数执行时都会被赋予一个非零且有用的序列号seq=<N>。这些函数可能会像正向传播中的原始函数一样,在稍后的二次反向传播中创建要执行的Function对象。反向传播和二次反向传播之间的关系在概念上与前向传播和反向传播之间的关系相同:函数仍然会发出带有当前序列号标签的数据范围,它们创建的Function对象仍然会存储这些序列号,并且在最终进行二次反向传播时,Function对象的apply()方法调用会被标记为stashed seq号码,这可以与来自第一次反向传播的seq号码进行比较。

torch.autograd.profiler.emit_itt(enabled=True, record_shapes=False)[源代码]

一个上下文管理器,使得每次自动微分操作都会发出一个ITT范围。

在使用 Intel(R) VTune 分析器运行程序时非常有用:

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

仪器和跟踪技术(ITT)API 允许你的应用程序在执行过程中使用不同的 Intel 工具生成和控制追踪数据的收集。这个上下文管理器用于标注 Intel VTune 分析器中的跟踪信息。借助此上下文管理器,你将在 Intel VTune 分析器 GUI 中看到带有标签的范围。

参数
  • enabled (bool, 可选) – 当设置为 enabled=False 时,此上下文管理器将不执行任何操作。默认值:True

  • record_shapes (bool, 可选) – 如果设置为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跟踪文件,并解析其中的自动微分注释。

调试与异常检测

torch.autograd.detect_anomaly(check_nan=True)[源代码]

一个启用自动微分引擎异常检测的上下文管理器。

这完成了两个任务:

  • 启用检测后运行前向传递,可以让反向传递打印出导致失败的反向函数的前向操作的追踪信息。

  • 如果 check_nanTrue,任何产生“nan”值的反向计算将触发错误。默认情况下,check_nanTrue

警告

此模式仅应在调试时启用,因为它会使程序的执行速度变慢。

示例

>>> 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
函数torch.autograd.set_detect_anomaly(mode, check_nan=True)[源代码]

一个上下文管理器,用于控制自动微分引擎的异常检测的开启和关闭。

set_detect_anomaly 根据参数 mode 来启用或禁用自动求导异常检测功能。它既可以作为上下文管理器使用,也可以直接作为函数调用。

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

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

  • check_nan (bool) – 是否在反向传播过程中生成“nan”值时触发异常

grad_mode.set_multithreading_enabled

一个上下文管理器,用于控制多线程反向传播的开启或关闭。

自动求导图

Autograd 提供了方法,让用户可以在反向传播过程中检查计算图并添加自定义行为。

torch.Tensorgrad_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

更新自动微分元数据,以跟踪给定张量是否被就地修改。

某些操作需要在前向传递过程中保存中间结果,以便进行后向传递。这些中间结果以属性的形式保存在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上,但在训练时不行,则此方法特别有用。另请参阅保存张量的钩子

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 返回的packed 对象作为参数,并返回一个与原始张量(传递给相应 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>)

警告

在任一边界的输入上执行原地操作可能导致未定义的行为。

警告

一次只能使用一组钩子。当递归嵌套此上下文管理器时,只会应用最内层的那组钩子。

classtorch.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
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
torch.autograd.graph.register_multi_grad_hook(tensors, fn, *, mode='all')[源代码]

注册一个多梯度反向挂钩。

有两种受支持的模式: "all""any"

"all"模式下,当tensors中每个张量的梯度计算完成后才会调用钩子。如果某个张量属于tensors但不在图中,或者对于当前.backward().grad()调用指定的任何输入不需要该张量来计算其梯度,则会忽略此张量,并且钩子不会等待其梯度被计算。

在每个非忽略张量的梯度计算完成后,将会调用 fn 并传入这些梯度。对于未计算梯度的张量,则会传递None

"any"模式下,当tensors中的某个张量计算出第一个梯度后,会触发钩子函数,并将这个梯度作为参数传递给钩子。

钩子不应修改其参数。

此函数返回一个包含 handle.remove() 方法的句柄,该方法用于移除钩子。

注意

关于此钩子的执行时间和与其他钩子的执行顺序,请参阅Backward Hooks 执行

示例:

>>> 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>对象,用于保存当前上下文管理器管理的状态,在调试过程中可能会非常有用。当你退出程序时,这些状态信息会被自动清理。

返回类型

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

本页目录