自動微分套件 - torch.autograd

torch.autograd 提供實作任意純量值函數自動微分的類別和函數。 它只需要對現有程式碼進行最少的更改 - 您只需要使用 requires_grad=True 關鍵字宣告應計算梯度的 Tensor。 目前，我們僅支援浮點數 Tensor 類型（半精度、單精度、雙精度和 bfloat16）和複數 Tensor 類型（cfloat、cdouble）的 Autograd。

backward	計算給定張量相對於圖形葉節點的梯度總和。
grad	計算並返回輸出相對於輸入的梯度總和。

正向模式自動微分

警告

此 API 處於 Beta 測試階段。 儘管函數簽章不太可能更改，但在我們認為它穩定之前，計畫改進運算子涵蓋範圍。

請參閱正向模式 AD 教學課程，以瞭解有關如何使用此 API 的詳細步驟。

forward_ad.dual_level	正向 AD 的上下文管理器，其中所有正向 AD 計算都必須在 dual_level 上下文中進行。
forward_ad.make_dual	將張量值与其切線相關聯，以建立用於正向 AD 梯度計算的「對偶張量」。
forward_ad.unpack_dual	解開「對偶張量」以取得其張量值和正向 AD 梯度。
forward_ad.enter_dual_level	進入新的正向梯度級別。
forward_ad.exit_dual_level	退出正向梯度級別。
forward_ad.UnpackedDualTensor	unpack_dual() 返回的 Namedtuple，其中包含對偶張量的原始組件和切線組件。

函數式高階 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 和推論模式之間差異的資訊，以及其他可能與這兩者混淆的相關機制。另請參閱局部停用梯度計算，以取得可用於局部停用梯度的函數清單。

預設梯度佈局

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

如果 param.grad 最初為 None

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

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

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

3. 如果 create_graph=False，backward() 會累積到 .grad 中，保留其步幅。

4. 如果 create_graph=True，backward() 會將 .grad 替換為新的張量 .grad + new grad，它會嘗試（但不保證）匹配現有的 .grad 的步幅。

建議使用預設行為（在第一次 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 的步幅，請在第一次 backward() 之前將 param.grad = 分配給具有所需步幅的歸零張量，並且永遠不要將其重設為 None。只要 create_graph=False，3 就可以確保您的佈局得到保留。4 表示即使 create_graph=True，您的佈局也*可能*得到保留。

張量的就地操作

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

就地正確性檢查

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

變數（已棄用）

警告

變數 API 已被棄用：不再需要變數來將自動梯度與張量一起使用。自動梯度自動支援 requires_grad 設定為 True 的張量。以下是有關已更改內容的快速指南

• Variable(tensor) 和 Variable(tensor, requires_grad) 仍然可以按預期工作，但它們返回的是張量而不是變數。

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

張量自動梯度函數

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。

函數

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	定義自定義自動梯度函數的正向傳播。
Function.backward	使用反向模式自動微分定義一個公式來微分操作。
Function.jvp	使用正向模式自動微分定義一個公式來微分操作。
Function.vmap	在 torch.vmap() 下定義此 autograd.Function 的行為。

上下文方法混合

建立新的 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	針對 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, use_mtia=False, experimental_config=None)

用於管理 Autograd 效能分析器狀態並保存結果摘要的上下文管理器。

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

參數

• enabled (bool, optional) – 將此設定為 False 會使此上下文管理器變成無效操作。

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

• use_device (str, optional) – 啟用裝置事件的計時。在使用 CUDA 時，會為每個張量運算增加大約 4 微秒的開銷。有效的裝置選項為「cuda」、「xpu」和「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，其中包含一個 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 等效能分析器程式庫使用的一組實驗性選項。請注意，不保證向後相容性。

範例

>>> 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 操作的 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 都會遞增。因此，與每個前向函數範圍關聯的 seq=<N> 註解告訴您，如果此前向函數創建了一個反向函數物件，則該反向物件將接收序號 N。在反向傳遞期間，封裝每個 C++ 反向函數的 apply() 呼叫的頂層範圍都使用 stashed seq=<M> 裝飾。 M 是創建反向物件時的序號。透過比較反向傳遞中的 stashed seq 編號和前向傳遞中的 seq 編號，您可以追蹤哪個前向運算創建了每個反向函數。

在反向傳遞期間執行的任何函數也都使用 seq=<N> 裝飾。在預設的反向傳遞期間（使用 create_graph=False），此資訊無關緊要，事實上，對於所有此類函數，N 可能只是 0。只有與反向函數物件的 apply() 方法相關聯的頂層範圍才有用，作為將這些函數物件與之前的前向傳遞關聯起來的一種方式。

雙重重播

另一方面，如果正在進行使用 create_graph=True 的反向傳遞（換句話說，如果您正在設置雙重重播），則在反向傳遞期間，每個函數的執行都會給予一個非零且有用的 seq=<N>。這些函數本身可能會創建函數物件，以便稍後在雙重重播期間執行，就像前向傳遞中的原始函數一樣。反向傳遞和雙重重播之間的關係在概念上與前向傳遞和反向傳遞之間的關係相同：函數仍然發出帶有當前序號標籤的範圍，它們創建的函數物件仍然儲存這些序號，並且在最終的雙重重播期間，函數物件的 apply() 範圍仍然使用 stashed seq 編號進行標記，這些編號可以與反向傳遞中的 seq 編號進行比較。

class torch.autograd.profiler.emit_itt(enabled=True, record_shapes=False)

上下文管理器，可讓每個自動微分操作發出 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，則封裝每個自動微分操作的 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 追蹤文件並解析自動微分註解。

除錯和異常檢測

class torch.autograd.detect_anomaly(check_nan=True)

上下文管理器，用於為自動微分引擎啟用異常檢測。

這會做兩件事

• 在啟用檢測的情況下運行前向傳遞將允許反向傳遞打印創建失敗的反向函數的前向操作的追蹤堆棧。

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

上下文管理器，用於啟用或關閉自動微分引擎的異常檢測。

set_detect_anomaly 將根據其參數 mode 啟用或關閉自動微分異常檢測。它可以用作上下文管理器或函數。

有關異常檢測行為的詳細信息，請參見上面的 detect_anomaly。

參數

• mode (bool) – 標記是否啟用異常檢測（True）或關閉（False）。

• check_nan (bool) – 標記在反向傳遞產生「nan」時是否引發錯誤

grad_mode.set_multithreading_enabled

上下文管理器，用於啟用或關閉多線程反向傳遞。

自動微分圖

自動微分公開了一些方法，允許用戶檢查圖並在反向傳遞期間插入行為。

如果 torch.Tensor 是由自動微分記錄的操作的輸出（即，啟用了 grad_mode 並且至少有一個輸入需要梯度），則 grad_fn 屬性將保存一個 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	更新自動微分元數據，追蹤給定的張量是否被就地修改。

某些操作需要在前向傳遞期間保存中間結果，以便執行反向傳遞。這些中間結果將作為屬性保存在 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)

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

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

在這種情況下，每次操作保存用於反向傳遞的張量時，都會調用 pack_hook 函數（這包括使用 save_for_backward() 保存的中間結果，以及由 PyTorch 定義的操作記錄的結果）。然後，pack_hook 的輸出將存儲在計算圖中，而不是原始張量。

當需要訪問保存的張量時（即執行 torch.Tensor.backward() 或 torch.autograd.grad() 時），將調用 unpack_hook。它將 pack_hook 返回的*打包*對象作為參數，並且應該返回一個與原始張量（作為輸入傳遞給相應的 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
>>>
>>> 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, 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

class torch.autograd.graph.disable_saved_tensors_hooks(error_message)

上下文管理器，用於停用已儲存張量的預設鉤子功能。

如果您要建立的功能不適用於已儲存張量的預設鉤子，則此功能很有用。

參數

error_message (str) – 當已儲存張量的預設鉤子在被停用時被使用時，會引發帶有此錯誤訊息的 RuntimeError。

範例

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

註冊一個多梯度反向鉤子。

有兩種支援的模式："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

允許變更已儲存用於反向傳播的張量的上下文管理器。

在此上下文管理器下，已儲存用於反向傳播的張量在變更時會被複製，因此原始版本仍然可以在反向傳播期間使用。通常，變更已儲存用於反向傳播的張量將導致在反向傳播期間使用它時引發錯誤。

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

回傳值

一個 _AllowMutationOnSavedContext 物件，儲存此上下文管理器管理的狀態。此物件對於除錯很有用。上下文管理器管理的狀態將在退出時自動清除。

範例

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