• Docs >
  • Automatic differentiation package - torch.autograd
Shortcuts

自动微分包 - torch.autograd

torch.autograd 提供了实现任意标量值函数的自动微分的类和函数。它对现有代码的改动最小 - 您只需要使用 requires_grad=True 关键字声明需要计算梯度的 Tensor。目前,我们仅支持浮点型 Tensor 类型(半精度、单精度、双精度和 bfloat16)和复数 Tensor 类型(cfloat、cdouble)的自动微分。

backward

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

grad

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

前向模式自动微分

警告

此API处于测试阶段。尽管函数签名几乎不可能更改,但我们计划在考虑将其视为稳定版本之前改进操作符覆盖范围。

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

forward_ad.dual_level

用于前向自动微分的上下文管理器,所有前向自动微分计算必须在 dual_level 上下文中进行。

forward_ad.make_dual

将张量值与其切线关联,以创建用于前向AD梯度计算的“对偶张量”。

forward_ad.unpack_dual

解包一个“双重张量”以获取其张量值和前向自动微分梯度。

forward_ad.enter_dual_level

输入一个新的前向梯度级别。

forward_ad.exit_dual_level

退出前向梯度级别。

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

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

functional.vjp

计算向量 v 与给定函数在输入点处的雅可比矩阵的点积。

functional.jvp

计算给定函数在输入点处的雅可比矩阵与向量v的点积。

functional.vhp

计算向量v与给定标量函数在指定点的Hessian矩阵的点积。

functional.hvp

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

本地禁用梯度计算

有关no-grad和推理模式之间的区别以及可能与两者混淆的其他相关机制的更多信息,请参阅局部禁用梯度计算。另请参阅局部禁用梯度计算,以获取可用于局部禁用梯度的函数列表。

默认渐变布局

当一个非稀疏的 param 在执行 torch.autograd.backward()torch.Tensor.backward() 时接收到一个非稀疏的梯度,param.grad 会按照如下方式累积。

如果 param.grad 最初是 None

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

  2. 否则,.grad 是以行优先连续步长创建的。

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

  1. 如果 create_graph=Falsebackward() 会就地累加到 .grad 中,从而保留其步幅。

  2. 如果 create_graph=Truebackward() 会将 .grad 替换为一个新张量 .grad + new grad,这会尝试(但不保证)匹配预先存在的 .grad 的步幅。

默认行为(在第一次调用 backward() 之前让 .gradNone,以便它们的布局根据 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都会跟踪应用于它们的就地操作,如果实现检测到某个张量在一个函数中被保存用于反向传播,但随后被就地修改,一旦开始反向传播,就会引发错误。这确保了如果您使用就地函数且没有看到任何错误,您可以确定计算的梯度是正确的。

变量(已弃用)

警告

Variable API 已被弃用:使用 autograd 处理张量时不再需要 Variables。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)

张量自动求导函数

torch.Tensor.grad

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

torch.Tensor.requires_grad

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

torch.Tensor.is_leaf

所有requires_gradFalse的张量将按惯例成为叶子张量。

torch.Tensor.backward([梯度, ...])

计算当前张量相对于图叶的梯度。

torch.Tensor.detach

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

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
>>>
>>> # 通过调用apply方法来使用它:
>>> output = Exp.apply(input)

Function.forward

定义自定义 autograd 函数的 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

设置是否具体化梯度张量。

自定义函数工具

用于反向方法的装饰器。

function.once_differentiable

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

function.BackwardCFunction

此类用于内部自动梯度工作。

function.InplaceFunction

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

function.NestedIOFunction

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

数值梯度检查

gradcheck

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

gradgradcheck

通过小的有限差分计算的梯度与关于浮点型或复数型且requires_grad=True的张量的解析梯度进行比较,检查inputsgrad_outputs的梯度。

GradcheckError

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

分析器

Autograd 包含一个分析器,可以让你检查模型中不同操作符的成本——无论是在 CPU 还是 GPU 上。目前实现了三种模式——仅使用 profile 的 CPU 模式。基于 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, use_mtia=False, experimental_config=None)[源代码]

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

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

Parameters
  • 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 (布尔值, 可选) – 跟踪张量内存的分配/释放。

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

  • with_modules (bool) – 记录与操作的调用栈相对应的模块层次结构(包括函数名称)。例如,如果模块 A 的 forward 调用了模块 B 的 forward,其中包含一个 aten::add 操作,那么 aten::add 的模块层次结构是 A.B。请注意,目前此支持仅存在于 TorchScript 模型中,而不适用于急切模式模型。

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

  • use_cpu (布尔值, 可选) – 分析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):  # 任何正常的 Python 代码,真的!
>>>         y = x ** 2
>>>         y.backward()
>>> # 注意:为了简洁起见,删除了一些列
>>> print(prof.key_averages().table(sort_by="self_cpu_time_total"))
-----------------------------------  ---------------  ---------------  ---------------
名称                                 总 CPU 时间       平均 CPU 时间    调用次数
-----------------------------------  ---------------  ---------------  ---------------
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 -- <常规命令在这里>

不幸的是,没有方法强制 nvprof 将收集的数据刷新到磁盘,因此对于 CUDA 分析,必须使用此上下文管理器来注释 nvprof 跟踪并在检查它们之前等待进程退出。然后,可以使用 NVIDIA Visual Profiler (nvvp) 来可视化时间线,或者使用 torch.autograd.profiler.load_nvprof() 加载结果以进行检查,例如在 Python REPL 中。

Parameters
  • enabled (bool, 可选) – 设置 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)  # 预热CUDA内存分配器和分析器
...     with torch.autograd.profiler.emit_nvtx():
...         model(x)

前向-后向相关性

在Nvidia Visual Profiler中查看使用emit_nvtx创建的配置文件时, 将每个反向传播操作与其对应的前向传播操作相关联可能会很困难。 为了简化此任务,emit_nvtx将其生成的范围附加序列号信息。

在前向传播过程中,每个函数范围都带有 seq= 的装饰。seq 是一个运行中的计数器,每当创建并存储一个新的反向传播函数对象时,它就会递增。因此,与每个前向传播函数范围相关联的 seq= 注释告诉你,如果这个前向传播函数创建了一个反向传播函数对象,那么该反向对象将接收到序列号 N。在反向传播过程中,每个 C++ 反向传播函数的 apply() 调用都被带有 stashed seq= 的顶级范围所包围。M 是反向对象创建时所带的序列号。通过比较反向传播中的 stashed seq 号与前向传播中的 seq 号,你可以追踪到哪个前向操作创建了每个反向传播函数。

在反向传播过程中执行的任何函数也会被装饰为 seq=。在默认的反向传播(使用 create_graph=False)中,这些信息是无关紧要的,实际上,对于所有这些函数,N 可能只是 0。只有与反向传播函数对象的 apply() 方法相关联的顶级范围是有用的,因为它们可以用来将这些函数对象与之前的正向传播过程相关联。

反向传播

另一方面,如果正在进行一个带有 create_graph=True 的反向传播(换句话说,如果你正在为双重反向传播做准备),反向传播期间的每个函数的执行都会被赋予一个非零的、有用的 seq=。这些函数本身可能会创建 Function 对象,以便在双重反向传播期间稍后执行,就像前向传播中的原始函数一样。反向传播与双重反向传播之间的关系在概念上与前向传播与反向传播之间的关系相同:函数仍然发出带有当前序列号标记的范围,它们创建的 Function 对象仍然存储这些序列号,并且在最终的双重反向传播期间,Function 对象的 apply() 范围仍然带有 stashed seq 编号,这些编号可以与反向传播中的 seq 编号进行比较。

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

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

在Intel(R) VTune Profiler下运行程序时非常有用:

vtune <--vtune-flags> <常规命令 在这里>

仪器和跟踪技术(ITT)API使您的应用程序能够在其执行过程中生成和控制跨不同Intel工具的跟踪数据的收集。 此上下文管理器用于注释Intel(R) VTune Profiling跟踪。借助此上下文管理器,您将能够在Intel(R) VTune Profiler GUI中看到标记的范围。

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

  • record_shapes (bool, 可选) – 如果 record_shapes=True,则每个自动求导操作的 itt 范围将附加有关该操作接收到的 Tensor 参数大小的信息,格式如下: [[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...] 非 Tensor 参数将表示为 []。 参数将按照后端操作接收它们的顺序列出。 请注意,此顺序可能与在 Python 端传递这些参数的顺序不匹配。另请注意,形状记录可能会增加 itt 范围创建的开销。 默认值:False

示例

>>> 使用 torch.autograd.profiler.emit_itt():
...     model(x)

profiler.load_nvprof

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

调试和异常检测

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

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

这做了两件事:

  • 启用检测运行前向传播将允许后向传播打印创建失败后向函数的操作的回溯。

  • 如果 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):
...         # 反向传播过程中出错
...         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
class torch.autograd.set_detect_anomaly(mode, check_nan=True)[源代码]

上下文管理器,用于设置自动求导引擎的异常检测开启或关闭。

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

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

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

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

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

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

在这种情况下,每次操作保存一个张量以供反向传播时(这包括使用save_for_backward()保存的中间结果,还包括由PyTorch定义的操作记录的张量),都会调用pack_hook函数。然后,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("打包", x)
...     return x
>>>
>>> def unpack_hook(x):
...     print("解包", 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
打包 tensor([1., 1., 1., 1., 1.], requires_grad=True)
打包 tensor([2., 2., 2., 2., 2.], grad_fn=)
>>> y.sum().backward()
解包 tensor([1., 1., 1., 1., 1.], requires_grad=True)
解包 tensor([2., 2., 2., 2., 2.], grad_fn=)

警告

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

警告

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

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

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

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

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

Parameters

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 和 b 保存在 GPU 上
...     with torch.autograd.graph.save_on_cpu():
...         prod_2 = prod_1 * c  # prod_1 和 c 保存在 CPU 上
...     y = prod_2 * a           # prod_2 和 a 保存在 GPU 上
...     return y
>>>
>>> y = f(a, b, c)
>>> del a, b, c  # 仅用于说明
>>> # a、b 和 prod_2 的内容仍然在 GPU 上存活
>>> # prod_1 和 c 的内容仅在 CPU 上存活
>>> y.sum().backward()  # 所有 CPU 张量被移回 GPU,用于反向传播
>>> # 所有中间张量在调用 backward 后被释放(删除)
class torch.autograd.graph.disable_saved_tensors_hooks(error_message)[源代码]

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

适用于如果你正在创建一个不使用保存张量默认钩子的功能。

Parameters

error_message (str) – 当保存的张量默认钩子被禁用时,会引发一个带有此错误消息的RuntimeError。

示例:

>>> message = "saved tensors default hooks are disabled"
>>> with torch.autograd.graph.disable_saved_tensors_hooks(message):
...     # 引发 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')[源代码]

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

有两种支持的模式:"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.allow_mutation_on_saved_tensors[源代码]

允许在此上下文管理器下对用于反向传播的张量进行修改。

在这个上下文管理器下,为反向传播保存的张量在发生变异时会被克隆, 因此原始版本仍然可以在反向传播期间使用。通常情况下,变异一个为反向传播保存的张量 会导致在反向传播期间使用时引发错误。

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

Returns

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

示例:

>>> import torch
>>> with torch.autograd.graph.allow_mutation_on_saved_tensors():
...     # 前向传播
...     a = torch.ones(2, 3, requires_grad=True)
...     b = a.clone()
...     out = (b**2).sum()
...     b.sin_()
...     # 反向传播
...     out.sum().backward()
...
tensor([[0.8415, 0.8415, 0.8415],
        [0.8415, 0.8415, 0.8415]], grad_fn=)
class torch.autograd.graph.GradientEdge(node, output_nr)

表示 autograd 图中给定梯度边缘的对象。 要获取将计算给定张量梯度的梯度边缘, 可以执行 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)) 是等价的。

优云智算