自動微分套件 - torch.autograd

torch.autograd 提供類別和函式，用於實作任意純量值函式的自動微分。

它只需要對現有程式碼進行最少的修改 - 您只需要宣告 Tensor，並使用 requires_grad=True 關鍵字來指定要計算其梯度。目前，我們僅支援浮點數 Tensor 類型（半精度、浮點數、雙精度和 bfloat16）和複數 Tensor 類型（cfloat、cdouble）的 autograd。

backward	計算給定 tensors 相對於圖葉的梯度總和。
grad	計算並傳回輸出相對於輸入的梯度總和。

前向模式自動微分

警告

此 API 處於 beta 階段。 即使函式簽章不太可能變更，我們也計畫在將其視為穩定版本之前，改善運算子涵蓋範圍。

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

forward_ad.dual_level	用於前向 AD 的情境管理器 (context-manager)，所有前向 AD 計算都必須在 dual_level 情境中進行。
forward_ad.make_dual	將張量值與其切線 (tangent) 相關聯，以建立用於前向 AD 梯度計算的「雙重張量」(dual tensor)。
forward_ad.unpack_dual	解包 (unpack) 一個「雙重張量」，以取得其張量值及其前向 AD 梯度。
forward_ad.enter_dual_level	進入新的前向梯度層級 (forward grad level)。
forward_ad.exit_dual_level	退出前向梯度層級。
forward_ad.UnpackedDualTensor	由 unpack_dual() 傳回的 Namedtuple，包含雙重張量的原始 (primal) 和切線 (tangent) 成分。

函數式高階 API

警告

此 API 處於 beta 階段。 即使函數簽名不太可能改變，但在我們認為它穩定之前，還計劃對性能進行重大改進。

本節包含基於上述基本 API 建構的 autograd 的高階 API，允許您計算 Jacobian 矩陣、Hessian 矩陣等。

此 API 適用於使用者提供的函數，這些函數僅將張量 (Tensor) 作為輸入並僅返回張量。如果您的函數接受的其他參數不是張量，或者是不需要梯度 (requires_grad) 的張量，您可以使用 lambda 來捕獲它們。例如，對於一個函數 f，它接受三個輸入，一個是我們要計算 Jacobian 矩陣的張量，另一個應該被視為常數的張量，以及一個布林標誌，如 f(input, constant, flag=flag)，你可以這樣使用它：functional.jacobian(lambda x: f(x, constant, flag=flag), input)。

functional.jacobian	計算給定函數的 Jacobian 矩陣。
functional.hessian	計算給定純量函數的 Hessian 矩陣。
functional.vjp	計算向量 v 與給定函數在輸入給定點的 Jacobian 矩陣之間的點積。
functional.jvp	計算給定函數在輸入給定點的 Jacobian 矩陣與向量 v 之間的點積。
functional.vhp	計算向量 v 與給定純量函數在指定點的 Hessian 矩陣之間的點積。
functional.hvp	計算純量函數的 Hessian 矩陣與向量 v 在指定點的點積。

局部禁用梯度計算

有關 no-grad 和 inference 模式之間差異以及其他可能與這兩者混淆的相關機制的更多資訊，請參閱局部禁用梯度計算。另請參閱 局部禁用梯度計算，以取得可用於局部禁用梯度的函數列表。

預設梯度佈局

當一個非稀疏的 param 在 torch.autograd.backward() 或 torch.Tensor.backward() 期間收到一個非稀疏梯度時，param.grad 會按照以下方式累積。

如果 param.grad 最初是 None

1. 如果 param 的記憶體是非重疊且密集的，則 .grad 會以符合 param 的步幅 (strides) (因此符合 param 的佈局) 來建立。

2. 否則，.grad 會以 rowmajor-contiguous 步幅建立。

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

3. 如果 create_graph=False，backward() 會就地 (in-place) 累積到 .grad 中，這會保留其步幅。

4. 如果 create_graph=True，backward() 會將 .grad 替換為一個新的張量 .grad + new grad，它會嘗試 (但不保證) 符合原有的 .grad 的步幅。

建議使用預設行為 (讓 .grad 在第一次 backward() 之前為 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，您的佈局也可能會被保留。

張量的就地操作

在 autograd 中支援就地操作是一件困難的事情，我們不鼓勵在大多數情況下使用它們。 Autograd 的積極緩衝區釋放和重用使其非常有效，並且很少有就地操作實際上能顯著降低記憶體使用量。 除非您在繁重的記憶體壓力下操作，否則您可能永遠不需要使用它們。

就地正確性檢查

所有 Tensor 都會追蹤應用於它們的就地操作，並且如果實現檢測到張量被保存在其中一個函數中以用於 backward，但在之後被就地修改，則一旦開始 backward 傳遞，就會引發錯誤。 這確保了如果您使用就地函數且沒有看到任何錯誤，您可以確定計算出的梯度是正確的。

Variable (已棄用)

警告

Variable API 已被棄用：不再需要使用 Variable 即可將 autograd 與 tensors 一起使用。Autograd 會自動支援 requires_grad 設定為 True 的 Tensors。以下快速指南說明了哪些已變更

• Variable(tensor) 和 Variable(tensor, requires_grad) 仍如預期運作，但它們會傳回 Tensors 而非 Variables。

• var.data 與 tensor.data 相同。

• 諸如 var.backward(), var.detach(), var.register_hook() 等方法現在可以使用相同的方法名稱在 tensors 上運作。

此外，現在可以使用工廠方法（例如 torch.randn()、torch.zeros()、torch.ones() 等）建立 requires_grad=True 的 tensors，如下所示

autograd_tensor = torch.randn((2, 3, 4), requires_grad=True)

Tensor autograd 函數

torch.Tensor.grad	預設情況下，此屬性為 None，並且在第一次呼叫 backward() 計算 self 的梯度時變為 Tensor。
torch.Tensor.requires_grad	如果需要計算此 Tensor 的梯度，則為 True，否則為 False。
torch.Tensor.is_leaf	所有具有 requires_grad 且為 False 的 Tensors 依照慣例將會是葉節點 Tensors。
torch.Tensor.backward([gradient, ...])	計算當前 tensor 相對於圖形葉節點的梯度。
torch.Tensor.detach	傳回與目前圖形分離的新 Tensor。
torch.Tensor.detach_	將 Tensor 從建立它的圖形中分離，使其成為葉節點。
torch.Tensor.register_hook(hook)	註冊一個向後 hook。
torch.Tensor.register_post_accumulate_grad_hook(hook)	註冊一個在梯度累積後執行的向後 hook。
torch.Tensor.retain_grad()	使此 Tensor 能夠在 backward() 期間填入其 grad。

Function

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
>>>
>>> # Use it by calling the apply method:
>>> output = Exp.apply(input)

Function.forward	定義自訂 autograd Function 的前向傳遞。
Function.backward	定義使用向後模式自動微分來微分運算的公式。
Function.jvp	定義使用前向模式自動微分來微分運算的公式。
Function.vmap	定義 torch.vmap() 下此 autograd.Function 的行為。

內容方法混合

建立新的 Function 時，以下方法可用於 ctx。

function.FunctionCtx.mark_dirty	將指定的 tensors 標記為在原地操作中已修改。
function.FunctionCtx.mark_non_differentiable	將輸出標記為不可微分。
function.FunctionCtx.save_for_backward	儲存給定的 tensors 以供將來呼叫 backward()。
function.FunctionCtx.set_materialize_grads	設定是否實現 grad tensors。

自訂 Function 實用程式

向後方法裝飾器。

function.once_differentiable

用於建立 PyTorch 實用程式的基底自訂 Function

function.BackwardCFunction	此類別用於內部 autograd 工作。
function.InplaceFunction	此類別僅用於向後相容性原因。
function.NestedIOFunction	此類別僅用於向後相容性原因。

數值梯度檢查

gradcheck	檢查透過小幅度有限差分計算的梯度，是否與 inputs 中具有浮點數或複數類型，且 requires_grad=True 的 tensors 的解析梯度一致。
gradgradcheck	檢查透過小幅度有限差分計算的梯度的梯度，是否與 inputs 和 grad_outputs 中具有浮點數或複數類型，且 requires_grad=True 的 tensors 的解析梯度一致。
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)

管理 autograd 分析器狀態並保存結果摘要的上下文管理器。

在底層，它只是記錄 C++ 中執行的函數的事件，並將這些事件公開給 Python。您可以將任何程式碼包裝在其中，它只會報告 PyTorch 函數的執行時間。注意：分析器是線程本地的，並且會自動傳播到異步任務中。

參數

• enabled (bool, optional) – 如果將此設定為 False，則此上下文管理器將不會執行任何操作。

• use_cuda (bool, optional) – 使用 cudaEvent API 啟用 CUDA 事件的計時。（將被棄用）

• use_device (str, optional) – 啟用裝置事件的計時。使用 cuda 時，每個 tensor 操作大約增加 4us 的開銷。有效的裝置選項為 ‘cuda’、‘xpu’、‘mtia’ 和 ‘privateuseone’。

• record_shapes (bool, optional) – 如果設定了形狀記錄，則會收集有關輸入維度的資訊。這允許人們查看在底層使用了哪些維度，並使用 prof.key_averages(group_by_input_shape=True) 依據這些維度進一步分組。請注意，形狀記錄可能會扭曲您的分析資料。建議使用單獨的運行（有和沒有形狀記錄）來驗證計時。對於最底層的事件（在巢狀函數呼叫的情況下），偏差很可能可以忽略不計。但對於較高層級的函數，由於形狀收集，總的 self cpu 時間可能會被人為地增加。

• with_flops (bool, optional) – 如果設定了 with_flops，分析器將使用運算子的輸入形狀來估算 FLOPs（浮點運算）值。這允許人們估算硬體效能。目前，此選項僅適用於矩陣乘法和 2D 卷積運算子。

• profile_memory (bool, optional) – 追蹤 tensor 記憶體分配/釋放。

• with_stack (bool, optional) – 記錄 ops 的原始碼資訊（檔案和行號）。

• with_modules (bool) – 記錄與 op 的呼叫堆疊相對應的模組層級結構（包括函數名稱）。例如，如果模組 A 的 forward 呼叫模組 B 的 forward，其中包含 aten::add op，則 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 累積

範例

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

上下文管理器，使每個 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 op 的 nvtx 範圍將附加有關該 op 接收的 Tensor 參數大小的資訊，格式如下： [[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...]。非 Tensor 參數將由 [] 表示。參數將按照後端 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)

正向-反向關聯

在使用 emit_nvtx 建立的 Nvidia Visual Profiler 中檢視分析時，將每個反向傳遞 op 與相應的正向傳遞 op 相關聯可能很困難。為了簡化這項任務，emit_nvtx 將序列號資訊附加到其產生的範圍。

在正向傳遞期間，每個函數範圍都標記有 seq=<N>。seq 是一個運行計數器，每次建立一個新的反向 Function 物件並隱藏以進行反向傳遞時，該計數器都會遞增。因此，與每個正向函數範圍關聯的 seq=<N> 註釋會告訴您，如果此正向函數建立了一個反向 Function 物件，則該反向物件將接收序列號 N。在反向傳遞期間，包裝每個 C++ 反向 Function 的頂層範圍 apply() 呼叫標記有 stashed seq=<M>。M 是建立反向物件時使用的序列號。通過比較反向傳遞中的 stashed seq 數字與正向傳遞中的 seq 數字，您可以追蹤哪個正向 op 建立了每個反向 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)

Context manager（上下文管理器），使每個 autograd 運算發出一個 ITT 範圍。

在 Intel(R) VTune Profiler 下運行程式時，它很有用。

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

Instrumentation and Tracing Technology (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 範圍將附加有關該運算接收到的 Tensor 參數大小的資訊，格式如下：[[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...]。非 Tensor 參數將以 [] 表示。參數將按照後端運算接收的順序列出。請注意，此順序可能與這些參數在 Python 端傳遞的順序不符。另請注意，形狀記錄可能會增加 itt 範圍建立的開銷。預設值：False

範例

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

profiler.load_nvprof

開啟一個 nvprof 追蹤檔案並解析 autograd 註解。

偵錯和異常檢測

class torch.autograd.detect_anomaly(check_nan=True)

Context-manager（上下文管理器），用於啟用 autograd 引擎的異常檢測。

這會做兩件事

• 在啟用檢測的情況下執行前向傳遞將允許後向傳遞列印出創建失敗的後向函數的前向運算的追蹤資訊。

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

Context-manager（上下文管理器），用於開啟或關閉 autograd 引擎的異常檢測。

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

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

參數

• mode (bool) – 標記是否啟用異常檢測 (True)，或停用 (False)。

• check_nan (bool) – 標記是否在後向傳播產生 “nan” 時引發錯誤

grad_mode.set_multithreading_enabled

Context-manager（上下文管理器），用於開啟或關閉多執行緒後向傳播。

Autograd 圖

Autograd 公開了一些方法，允許人們在後向傳遞期間檢查圖並介入行為。

如果張量是 autograd 記錄的運算的輸出（即，grad_mode 已啟用且至少一個輸入需要梯度），則 torch.Tensor 的 grad_fn 屬性保留一個 torch.autograd.graph.Node，否則保留 None。

graph.Node.name	傳回名稱。
graph.Node.metadata	傳回元資料。
graph.Node.next_functions
graph.Node.register_hook	註冊一個後向 hook。
graph.Node.register_prehook	註冊一個後向 pre-hook。
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

您還可以定義應如何使用 hook 來封裝/解封裝這些儲存的張量。常見的應用程式是通過將這些中間結果儲存到磁碟或 CPU 而不是將它們留在 GPU 上來換取記憶體以進行計算。如果您注意到您的模型在評估期間適合 GPU，但在訓練期間不適合，這尤其有用。另請參閱 儲存張量的 Hook。

class torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook)

Context-manager（上下文管理器），用於為儲存的張量設定一對封裝/解封裝 hook。

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

在此上下文中，每次操作儲存張量以進行反向傳播時（包括使用 save_for_backward() 儲存的中間結果，以及 PyTorch 定義的操作記錄的結果），都會呼叫 pack_hook 函數。pack_hook 的輸出會儲存在計算圖中，而不是原始張量。

當需要存取已儲存的張量時，即執行 torch.Tensor.backward() 或 torch.autograd.grad() 時，會呼叫 unpack_hook。它接受 pack_hook 返回的已包裝物件作為參數，並且應返回一個與原始張量（作為輸入傳遞給相應的 pack_hook）具有相同內容的張量。

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

警告

對任一 hook 的輸入執行 inplace 操作可能會導致未定義的行為。

警告

一次只允許一對 hook。當遞迴巢狀使用此上下文管理器時，只會應用最內層的一對 hook。

class torch.autograd.graph.save_on_cpu(pin_memory=False, device_type='cuda')

在此上下文管理器下，前向傳遞儲存的張量將儲存在 CPU 上，然後檢索以進行反向傳遞。

當在此上下文管理器中執行操作時，在前向傳遞期間儲存在圖中的中間結果將被移動到 CPU，然後在反向傳遞需要時複製回原始裝置。如果圖已在 CPU 上，則不會執行張量複製。

使用此上下文管理器來利用計算力換取 GPU 記憶體使用量（例如，當您的模型在訓練期間不適合 GPU 記憶體時）。

參數

pin_memory (bool) – 如果為 True，則張量將在包裝期間儲存到 CPU pinned memory，並在解包期間非同步複製到 GPU。預設值為 False。另請參閱 使用 pinned memory 緩衝區。

範例

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

禁用已儲存張量預設 hook 功能的上下文管理器。

如果您正在建立一個與已儲存張量預設 hook 不相容的功能，這會很有用。

參數

error_message (str) – 當已停用已儲存張量預設 hook 時，如果使用了這些 hook，則會引發一個包含此錯誤訊息的 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')

註冊一個 multi-grad backward hook。

支援兩種模式："all" 和 "any"。

在 "all" 模式下，在計算完 tensors 中每個張量的梯度後，將呼叫該 hook。 如果張量在 tensors 中，但不是圖的一部分，或者不需要張量來計算當前 .backward() 或 .grad() 呼叫指定的任何 inputs 的梯度，則將忽略此張量，並且 hook 將不會等待計算其梯度。

在計算出每個非忽略張量的梯度後，將使用這些梯度呼叫 fn。對於沒有計算其梯度的張量，將傳遞 None。

在 "any" 模式下，在計算出關於 tensors 中張量的第一個梯度後，將呼叫該 hook。將使用該梯度作為其參數呼叫該 hook。

該 hook 不應修改其參數。

此函數返回一個帶有 handle.remove() 方法的句柄，該方法刪除該 hook。

注意

有關何時執行此 hook 以及其執行順序相對於其他 hook 的更多資訊，請參閱 Backward Hooks execution。

範例

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

允許修改儲存以供反向傳播之張量的上下文管理器。

在此上下文管理器下，為了反向傳播而儲存的張量會在修改時被複製，因此原始版本在反向傳播期間仍然可以使用。正常情況下，修改為反向傳播儲存的張量將導致在反向傳播期間使用它時引發錯誤。

為了確保正確的行為，正向傳播和反向傳播都應在相同的上下文管理器下運行。

回傳

一個 `_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>)

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

返回類型

GradientEdge