自動微分包 - torch.autograd#

建立日期：2016 年 12 月 23 日 | 最後更新日期：2025 年 6 月 12 日

torch.autograd 提供了實現任意標量值函式自動微分的類和函式。

它只需要對現有程式碼進行最少的更改——您只需使用 requires_grad=True 關鍵字宣告需要計算梯度的 Tensor。目前，我們僅支援浮點 Tensor 型別（half, float, double 和 bfloat16）以及複數 Tensor 型別（cfloat, cdouble）的自動微分。

backward

計算給定張量相對於圖葉子的梯度之和。

grad

計算並返回輸出相對於輸入的梯度之和。

前向模式自動微分#

警告

此 API 處於 beta 階段。儘管函式簽名非常可能不會改變，但在我們認為其穩定之前，計劃對其運算元覆蓋範圍進行改進。

有關如何使用此 API 的詳細步驟，請參閱前向模式 AD 教程。

`forward_ad.dual_level`	前向 AD 的上下文管理器，所有前向 AD 計算都必須發生在 `dual_level` 上下文中。
`forward_ad.make_dual`	將張量值與其切線關聯起來，建立一個“對偶張量”用於前向 AD 梯度計算。
`forward_ad.unpack_dual`	解壓“對偶張量”以獲取其張量值和前向 AD 梯度。
`forward_ad.enter_dual_level`	進入一個新的前向 grad 級別。
`forward_ad.exit_dual_level`	退出前向 grad 級別。
`forward_ad.UnpackedDualTensor`	由 `unpack_dual()` 返回的命名元組，包含對偶張量的原始分量和切線分量。

功能性高階 API#

警告

此 API 處於 beta 階段。儘管函式簽名非常可能不會改變，但在我們認為其穩定之前，計劃對其進行重大的效能改進。

本節包含 autograd 的高階 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` 在指定點的點積。

區域性停用梯度計算#

有關 no-grad 和 inference 模式之間差異以及其他可能與之混淆的相關機制的更多資訊，請參閱本地停用梯度計算。有關用於本地停用梯度的函式列表，請參閱本地停用梯度計算。

預設梯度佈局#

當非稀疏 param 在 torch.autograd.backward() 或 torch.Tensor.backward() 中接收到非稀疏梯度時，param.grad 會按如下方式累加。

如果 param.grad 最初為 None

如果 param 的記憶體是非重疊且稠密的，則 .grad 會以與 param 匹配的步幅建立（從而匹配 param 的佈局）。
否則，.grad 會以行主序連續的步幅建立。

如果 param 已經有一個非稀疏的 .grad 屬性

如果 create_graph=False，backward() 會原地累加到 .grad 中，這會保留其步幅。
如果 create_graph=True，backward() 會將 .grad 替換為一個新張量 .grad + new grad，該新張量會嘗試（但不保證）匹配預先存在的 .grad 的步幅。

建議採用預設行為（在第一次 backward() 之前讓 .grads 為 None，以便根據 1 或 2 建立其佈局，並根據 3 或 4 隨時間保留）以獲得最佳效能。呼叫 model.zero_grad() 或 optimizer.zero_grad() 不會影響 .grad 佈局。

事實上，在每個累加階段之前將所有 .grads 重置為 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 都會跟蹤應用到它們的原地操作，如果實現檢測到一個張量在某個函式中為 backward 儲存，但之後又被原地修改，則在 backward 傳遞開始時會引發錯誤。這確保瞭如果您使用原地函式且沒有看到任何錯誤，您可以確信計算出的梯度是正確的。

Variable (已棄用)#

警告

Variable API 已棄用：使用 autograd 和張量不再需要 Variable。Autograd 自動支援 requires_grad 設定為 True 的張量。下面是一個關於變化之處的快速指南。

Variable(tensor) 和 Variable(tensor, requires_grad) 仍按預期工作，但它們返回張量而不是 Variable。
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)

張量 autograd 函式#

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

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

建立自定義 autograd.Function 的基類。

要建立自定義 autograd.Function，請繼承此類並實現 forward() 和 backward() 靜態方法。然後，要在前向傳播中使用自定義運算元，請呼叫類方法 apply。請勿直接呼叫 forward()。

為確保正確性和最佳效能，請確保您正在對 ctx 呼叫正確的方法，並使用 torch.autograd.gradcheck() 驗證您的 backward 函式。

有關如何使用此類，請參閱擴充套件 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`	定義此 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`	設定是否例項化 grad 張量。

自定義函式實用程式#

backward 方法的裝飾器。

function.once_differentiable

用於構建 PyTorch 實用程式的基類自定義 Function

`function.BackwardCFunction`	此類用於內部 autograd 工作。
`function.InplaceFunction`	此類僅出於向後相容性原因而存在。
`function.NestedIOFunction`	此類僅出於向後相容性原因而存在。

數值梯度檢查#

gradcheck

透過有限差分法計算的梯度與 inputs 中是浮點或複數型別且 requires_grad=True 的張量的分析梯度進行比較。

gradgradcheck

透過有限差分法計算的梯度與 inputs 和 grad_outputs 中是浮點或複數型別且 requires_grad=True 的張量的分析梯度進行比較。

GradcheckError

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, 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, custom_trace_id_callback=None)[source]#

上下文管理器，用於管理 autograd 分析器狀態並儲存結果摘要。

注意

這是後端，大多數人應該使用 torch.profiler 代替。

它只是在後臺記錄 C++ 中執行的函式事件，並將這些事件暴露給 Python。您可以將任何程式碼包裝其中，它只會報告 PyTorch 函式的執行時。注意：分析器是執行緒本地的，並且會自動傳播到非同步任務中。

引數

enabled (bool, optional) – 將此設定為 False 會使此上下文管理器無效（no-op）。
use_cuda (bool, optional) – 使用 cudaEvent API 啟用 CUDA 事件計時。（將棄用）
use_device (str, optional) – 啟用裝置事件計時。使用 cuda 時，每個張量操作會增加大約 4us 的開銷。有效裝置選項為“cuda”、“xpu”、“mtia”和“privateuseone”。
record_shapes (bool, optional) – 如果設定了形狀記錄，將收集輸入維度的資訊。這允許您檢視哪些維度在底層使用過，並使用 prof.key_averages(group_by_input_shape=True) 按它們進行分組。請注意，形狀記錄可能會歪曲您的分析資料。建議使用單獨的執行（帶形狀記錄和不帶形狀記錄）來驗證計時。最底層事件（在巢狀函式呼叫的情況下）的自 CPU 時間可能會因形狀收集而被人為增加。
with_flops (bool, optional) – 如果設定了 with_flops，分析器將使用運算元的輸入形狀估算 FLOPs（浮點運算）值。這允許您估算硬體效能。目前，此選項僅適用於矩陣乘法和二維卷積運算元。
profile_memory (bool, optional) – 跟蹤張量記憶體分配/去分配。
with_stack (bool, optional) – 記錄運算元的源資訊（檔名和行號）。
with_modules (bool) – 記錄與運算元呼叫棧對應的模組層次結構（包括函式名）。例如，如果模組 A 的 forward 呼叫了模組 B 的 forward，而模組 B 包含一個 aten::add 運算元，那麼 aten::add 的模組層次結構就是 A.B。請注意，目前此支援僅適用於 TorchScript 模型，而不適用於 eager 模式模型。
use_kineto (bool, optional) – 實驗性的，啟用 Kineto 分析器進行分析。
use_cpu (bool, optional) – 分析 CPU 事件；將其設定為 False 需要 use_kineto=True，並且可用於降低 GPU 僅分析的開銷。
experimental_config (_ExperimentalConfig) – Kineto 等分析器庫使用的實驗性選項集。注意，不保證向後相容性。
acc_events (bool) – 啟用跨多個分析週期累加 FunctionEvents。

警告

啟用記憶體分析或源屬性會產生額外的分析器開銷。

警告

不應遞迴呼叫此上下文管理器，即不允許巢狀例項。

警告

由於一些 CUDA 多程序限制（請參閱多程序中的 CUDA），不能使用 use_device = 'cuda' 的分析器來基準測試 num_workers > 0 的 DataLoaders。如果要基準測試資料載入，請使用 use_device = None 或 num_workers = 0。

示例

>>> 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`	用於訪問 interval 中 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, optional) – 將 enabled=False 設定為使此上下文管理器無效。預設值：True。
record_shapes (bool, optional) – 如果 record_shapes=True，則包裝每個 autograd 運算元的 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。只有與後向 Function 物件的 apply() 方法相關的頂級範圍才有意義，作為將這些 Function 物件與早期前向傳遞關聯的一種方式。

二次後向

另一方面，如果正在進行帶有 create_graph=True 的後向傳遞（換句話說，如果您正在為二次後向做準備），則後向期間執行的每個函式都會獲得一個非零的、有用的 seq=<N>。這些函式本身可能會建立 Function 物件，以便稍後在二次後向中執行，就像原始前向函式那樣。後向和二次後向之間的關係在概念上與前向和後向之間的關係相同：函式仍然發出當前序列號標記的範圍，它們建立的 Function 物件仍然會隱藏這些序列號，並且在最終的二次後向期間，Function 物件的 apply() 範圍仍會用 stashed seq 號進行標記，這些號可以與後向傳遞中的 seq 號進行比較。

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

上下文管理器，使每個 autograd 運算元都發出 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, optional) – 將 enabled=False 設定為使此上下文管理器無效。預設值：True。
record_shapes (bool, optional) – 如果 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]#

啟用 autograd engine 異常檢測的上下文管理器。

這做了兩件事：

啟用檢測後執行前向傳遞，將允許後向傳遞打印出建立失敗的後向函式的那個前向操作的堆疊跟蹤。
如果 check_nan 為 True，任何生成“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]#

根據其引數 mode 開啟或關閉 autograd engine 異常檢測的上下文管理器。

set_detect_anomaly 將根據其引數 mode 啟用或停用 autograd 異常檢測。它可用作上下文管理器或函式。

有關異常檢測行為的詳細資訊，請參閱上面的 detect_anomaly。

引數

mode (bool) – 標誌，指示是否啟用異常檢測（True）或停用（False）。
check_nan (bool) – 標誌，指示後向生成“nan”時是否引發錯誤。

grad_mode.set_multithreading_enabled

開啟或關閉多執行緒後向的上下文管理器。

Autograd 圖#

Autograd 提供的方法允許檢查圖並干預後向傳遞中的行為。

一個 torch.Tensor 的 grad_fn 屬性如果張量是 autograd 記錄的操作（即，grad_mode 已啟用且至少一個輸入需要梯度）的輸出，則包含一個 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`	更新 autograd 元資料，跟蹤給定張量是否被原地修改。

某些操作需要在前向傳遞期間儲存中間結果才能執行後向傳遞。這些中間結果儲存在 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]#

上下文管理器，用於為儲存的張量設定一對打包/解包鉤子。

使用此上下文管理器來定義操作的中間結果在儲存之前應如何打包，以及在檢索時如何解包。

在此上下文中，每當操作儲存一個張量用於後向（這包括使用 save_for_backward() 儲存的中間結果，以及 PyTorch 定義的操作記錄的中間結果）時，都會呼叫 pack_hook 函式。 pack_hook 的輸出隨後儲存在計算圖中，而不是原始張量。

unpack_hook 在需要訪問儲存的張量時被呼叫，即在執行 torch.Tensor.backward() 或 torch.autograd.grad() 時。它接受 pack_hook 返回的 *打包* 物件作為引數，並應返回一個在值、大小、dtype 和裝置上與原始張量（傳遞給相應的 pack_hook）相同的張量。

鉤子應具有以下簽名：

pack_hook(tensor: Tensor) -> Any

unpack_hook(Any) -> Tensor

其中 pack_hook 的返回值可作為 unpack_hook 的有效輸入。

一般來說，您希望 unpack_hook(pack_hook(t)) 在值、大小、dtype 和裝置方面等於 t。

示例

>>> def pack_hook(x):
...     print("Packing", x)
...     return x.detach()
>>>
>>> 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>)

警告

對任一鉤子的輸入執行原地操作可能會導致未定義的行為。

警告

一次只允許一對鉤子。在遞迴巢狀此上下文管理器時，僅應用最內層的一對鉤子。

警告

為避免引用迴圈，pack_hook 的返回值不能包含對輸入張量的引用。例如，使用 lambda x: x.detach() 而不是 lambda x: x 作為 pack 鉤子。

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。
返回型別: 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

class torch.autograd.graph.register_multi_grad_hook(tensors, fn, *, mode='all')[source]#

註冊多梯度後向鉤子。

支援兩種模式：“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]
>>>

返回型別: RemovableHandle

class torch.autograd.graph.allow_mutation_on_saved_tensors[source]#

允許在為後向儲存的張量上進行變異的上下文管理器。

在此上下文管理器下，為後向儲存的張量在變異時會被克隆，因此原始版本仍然可以在後向傳遞中使用。通常，變異一個為後向儲存的張量會在後向傳遞中使用它時引發錯誤。

為確保正確行為，前向和後向都應在相同的上下文管理器下執行。

返回: 一個 _AllowMutationOnSavedContext 物件，用於儲存此上下文管理器管理的 State。此物件可用於除錯目的。上下文管理器管理的 State 在退出時會自動清除。
返回型別: 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>)

class torch.autograd.graph.GradientEdge(node, output_nr, ownership_token=None)[source]#

表示 autograd 圖中給定梯度邊的物件。

要獲取將計算給定張量梯度的梯度邊，您可以執行 edge = autograd.graph.get_gradient_edge(tensor)。

torch.autograd.graph.get_gradient_edge(tensor)[source]#

獲取計算給定張量梯度的梯度邊。

特別是，呼叫 g = autograd.grad(loss, input) 和 g = autograd.grad(loss, get_gradient_edge(input)) 是等效的。

返回型別: GradientEdge

自動微分包 - torch.autograd#

前向模式自動微分#

功能性高階 API#

區域性停用梯度計算#

預設梯度佈局#

手動梯度佈局#

張量的原地操作#

原地正確性檢查#

Variable (已棄用)#

張量 autograd 函式#

Function#

上下文方法混合#

自定義函式實用程式#

數值梯度檢查#

分析器#

除錯和異常檢測#

Autograd 圖#

文件

教程

資源