自動微分包 - torch.autograd

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	進入新的前向梯度級別。
forward_ad.exit_dual_level	退出前向梯度級別。
forward_ad.UnpackedDualTensor	由 unpack_dual() 返回的 Namedtuple，包含對偶張量的原始分量和切線分量。

函式式高階 API

警告

此 API 處於 Beta 階段。儘管函式簽名極不可能改變，但在被認為是穩定版本之前，計劃進行重大的效能改進。

本節包含基於上述基本 API 構建的自動微分高階 API，允許您計算雅可比矩陣、海森矩陣等。

此 API 適用於使用者提供的函式，這些函式僅將 Tensor 作為輸入並僅返回 Tensor。如果您的函式接受非 Tensor 的其他引數或未設定 requires_grad 的 Tensor，您可以使用 lambda 表示式捕獲它們。例如，對於一個函式 f，它接受三個輸入：一個需要計算雅可比矩陣的 Tensor，一個應視為常量的 Tensor，以及一個布林標誌 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 的跨步是按行主序連續的。

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

3. 如果 create_graph=False，backward() 會就地累積到 .grad 中，從而保留其跨步。

4. 如果 create_graph=True，backward() 會用一個新的張量 .grad + new grad 替換 .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 都跟蹤應用到其上的就地操作，如果實現檢測到某個張量在一個函式中被儲存用於反向傳播，但之後被就地修改，一旦反向傳播開始，就會引發錯誤。這確保瞭如果您使用就地函式且沒有看到任何錯誤，您可以確信計算出的梯度是正確的。

Variable (已棄用)

警告

Variable API 已被棄用：使用張量進行自動微分不再需要 Variable。自動微分自動支援設定了 requires_grad 為 True 的 Tensor。下面是關於已更改內容的快速指南。

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

Tensor 自動微分函式

torch.Tensor.grad	此屬性預設為 None，並在第一次呼叫 backward() 計算 self 的梯度時變為 Tensor。
torch.Tensor.requires_grad	如果此 Tensor 需要計算梯度，則為 True，否則為 False。
torch.Tensor.is_leaf	按照慣例，所有 requires_grad 為 False 的 Tensor 都是葉節點 Tensor。
torch.Tensor.backward([gradient, ...])	計算當前張量相對於圖葉節點的梯度。
torch.Tensor.detach	返回一個新 Tensor，與當前計算圖分離。
torch.Tensor.detach_	將 Tensor 與建立它的計算圖分離，使其成為葉節點。
torch.Tensor.register_hook(hook)	註冊一個反向鉤子。
torch.Tensor.register_post_accumulate_grad_hook(hook)	註冊一個在梯度累積後執行的反向鉤子。
torch.Tensor.retain_grad()	使此 Tensor 的 grad 在 backward() 期間被填充。

Function

類 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 的前向計算。
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 工具類

用於 backward 方法的裝飾器。

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() 引發的錯誤。

分析器

自動微分包含一個分析器，允許您檢查模型中不同運算元的成本 - 包括在 CPU 和 GPU 上。目前實現了三種模式：僅限 CPU 使用 profile；基於 nvprof（同時記錄 CPU 和 GPU 活動）使用 emit_nvtx；以及基於 vtune 分析器使用 emit_itt。

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

管理自動微分分析器狀態並儲存結果摘要的上下文管理器。

在底層，它只是記錄 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, 可選) – 記錄運算元的源資訊（檔案和行號）。

• with_modules (bool) – 記錄與運算元呼叫棧對應的模組層次結構（包括函式名）。例如，如果模組 A 的 forward 呼叫了模組 B 的 forward，其中包含一個 aten::add 運算元，則 aten::add 的模組層次結構是 A.B。注意，目前此支援僅適用於 TorchScript 模型而非即時模式模型。

• use_kineto (bool, 可選) – 實驗性功能，啟用 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 會使此上下文管理器成為一個無操作（no-op）。預設值: True。

• record_shapes (bool, optional) – 如果 record_shapes=True，包裹每個 autograd 操作的 nvtx 範圍將附加關於該操作接收到的 Tensor 引數尺寸的資訊，格式如下： [[arg0.size(0), arg0.size(1), ...], [arg1.size(0), arg1.size(1), ...], ...] 非 Tensor 引數將用 [] 表示。引數將按照後端操作接收它們的順序進行列表。請注意，此順序可能與這些引數在 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)

上下文管理器，使每個 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 會使此上下文管理器成為一個無操作（no-op）。預設值: 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)

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

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

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 屬性持有一個 torch.autograd.graph.Node，如果該 tensor 是一個被 autograd 記錄的操作的輸出（即，autograd 模式已啟用並且至少有一個輸入需要梯度），否則為 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

您還可以使用鉤子定義這些儲存的 tensor 如何打包/解包。一個常見的應用是犧牲計算來換取記憶體，方法是將這些中間結果儲存到磁碟或 CPU，而不是將其留在 GPU 上。如果您的模型在評估期間適合 GPU，但在訓練期間不適合，這尤其有用。另請參閱 儲存的 tensor 的鉤子。

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

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

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

在此上下文中，`pack_hook` 函式將在每次一個操作為後向傳播儲存一個 tensor 時被呼叫（這包括使用 save_for_backward() 儲存的中間結果，以及由 PyTorch 定義的操作記錄的結果）。然後，`pack_hook` 的輸出被儲存在計算圖中，而不是原始 tensor。

當儲存的 tensor 需要被訪問時，即在執行 torch.Tensor.backward() 或 torch.autograd.grad() 時，`unpack_hook` 被呼叫。它以 `pack_hook` 返回的*打包*物件作為引數，並應返回一個 tensor，其內容與原始 tensor 相同（作為輸入傳遞給相應的 `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>)

警告

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

警告

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

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

上下文管理器，在此管理器下，前向傳播儲存的 tensor 將儲存在 CPU 上，然後為後向傳播進行檢索。

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

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

引數

pin_memory (bool) – 如果 True，tensor 將在打包期間儲存到 CPU 鎖頁記憶體（pinned memory），並在解包期間非同步複製到 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)

上下文管理器，停用儲存的 tensor 預設鉤子功能。

當您正在建立一個功能，但該功能與儲存的 tensor 預設鉤子不相容時很有用。

引數

error_message (str) – 當儲存的 tensor 預設鉤子已被停用但仍被使用時，將丟擲一個帶有此錯誤訊息的 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')

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

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

在 "all" 模式下，在計算完關於 tensors 中每個 tensor 的梯度後將呼叫該鉤子。如果一個 tensor 在 tensors 中但不屬於計算圖的一部分，或者如果對於當前 .backward() 或 .grad() 呼叫指定的任何 inputs 不需要一個 tensor 來計算梯度，這個 tensor 將被忽略，並且鉤子不會等待其梯度被計算。

在計算完所有非忽略 tensor 的梯度後，將使用這些梯度呼叫 fn。對於未計算其梯度的 tensor 將傳遞 None。

在 "any" 模式下，在計算完關於 tensors 中某個 tensor 的第一個梯度後將呼叫鉤子。將使用該梯度作為引數呼叫鉤子。

鉤子不應修改其引數。

此函式返回一個帶有方法 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

上下文管理器，在此管理器下允許修改為後向傳播儲存的 tensor。

在此上下文管理器下，為後向傳播儲存的 tensor 在修改時會被克隆，因此原始版本在後向傳播期間仍然可以使用。通常，修改為後向傳播儲存的 tensor 將導致在後向傳播期間使用它時丟擲錯誤。

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

返回

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

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

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

torch.autograd.graph.get_gradient_edge(tensor)

獲取用於計算給定 Tensor 梯度的梯度邊。

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

返回型別

GradientEdge