torch.utils.checkpoint

注意

檢查點 (Checkpointing) 的實作方式是在反向傳播期間，為每個檢查點區段重新執行一次正向傳播區段。這可能會導致持久狀態（例如 RNG 狀態）比沒有檢查點的情況更超前。預設情況下，檢查點包含調整 RNG 狀態的邏輯，以便使用 RNG（例如透過 dropout）的檢查點傳遞與非檢查點傳遞相比，具有確定性的輸出。根據檢查點操作的執行時間長短，儲存和恢復 RNG 狀態的邏輯可能會產生一定的效能損失。如果不需要與非檢查點傳遞相比的確定性輸出，請提供 preserve_rng_state=False 給 checkpoint 或 checkpoint_sequential 以省略在每次檢查點期間儲存和恢復 RNG 狀態。

儲存邏輯會儲存和恢復 CPU 和另一種裝置類型 (透過 _infer_device_type 從 Tensor 參數（不包括 CPU tensors）推斷裝置類型) 的 RNG 狀態到 run_fn。如果有多個裝置，裝置狀態只會為單一裝置類型的裝置儲存，其餘裝置將被忽略。因此，如果任何檢查點函數涉及隨機性，這可能會導致不正確的梯度。（請注意，如果 CUDA 裝置在檢測到的裝置中，它將被優先考慮；否則，將選擇遇到的第一個裝置。）如果沒有 CPU tensors，預設裝置類型狀態（預設值為 cuda，並且可以透過 DefaultDeviceType 設定為其他裝置）將被儲存和恢復。但是，該邏輯無法預測使用者是否會在 run_fn 本身內將 Tensors 移動到新裝置。因此，如果您在 run_fn 內將 Tensors 移動到新裝置（「新」表示不屬於 [目前裝置 + Tensor 參數的裝置] 的集合），則永遠無法保證與非檢查點傳遞相比的確定性輸出。

torch.utils.checkpoint.checkpoint(function, *args, use_reentrant=None, context_fn=<function noop_context_fn>, determinism_check='default', debug=False, **kwargs)

對模型或模型的一部分進行檢查點。

激活檢查點 (Activation checkpointing) 是一種以計算換取記憶體的技術。在檢查點區域的正向計算中，不保留反向傳播所需的 tensors，而是會重新計算它們，而不是將反向傳播所需的 tensors 保持到在反向傳播期間用於梯度計算。激活檢查點可以應用於模型的任何部分。

目前有兩種可用的檢查點實作方式，由 use_reentrant 參數決定。建議您使用 use_reentrant=False。請參考下面的註解，以討論它們的差異。

警告

如果在反向傳播期間 function 的呼叫與正向傳播不同，例如由於全域變數，則檢查點版本可能不等效，可能會導致引發錯誤或導致靜默地不正確的梯度。

警告

應明確傳遞 use_reentrant 參數。在 2.4 版中，如果未傳遞 use_reentrant，我們將引發例外。如果您使用的是 use_reentrant=True 變體，請參考下面的註解，以了解重要的考慮事項和潛在的限制。

注意

檢查點的 reentrant 變體 (use_reentrant=True) 和檢查點的 non-reentrant 變體 (use_reentrant=False) 在以下方面有所不同

• Non-reentrant 檢查點會在所有需要的中間激活被重新計算後立即停止重新計算。預設情況下啟用此功能，但可以使用 set_checkpoint_early_stop() 停用。Reentrant 檢查點始終在反向傳播期間完整地重新計算 function。

• Reentrant 變體不會在正向傳播期間記錄 autograd 圖，因為它在 torch.no_grad() 下執行正向傳播。Non-reentrant 版本確實會記錄 autograd 圖，允許您在檢查點區域內對圖執行反向傳播。

• Reentrant 檢查點僅支援不帶 inputs 參數的反向傳播的 torch.autograd.backward() API，而 non-reentrant 版本支援所有執行反向傳播的方式。

• 對於 reentrant 變體，至少一個輸入和輸出必須具有 requires_grad=True。如果未滿足此條件，則模型中經過檢查點的部分將沒有梯度。Non-reentrant 版本沒有此要求。

• Reentrant 版本不將巢狀結構（例如，自訂物件、列表、字典等）中的 tensors 視為參與 autograd，而 non-reentrant 版本會。

• Reentrant 檢查點不支援具有從計算圖中分離的 tensors 的檢查點區域，而 non-reentrant 版本支援。對於 reentrant 變體，如果檢查點區段包含使用 detach() 分離或具有 torch.no_grad() 的 tensors，則反向傳播將引發錯誤。這是因為 checkpoint 使所有輸出都需要梯度，並且當 tensor 被定義為在模型中沒有梯度時，這會導致問題。為了避免這種情況，請在 checkpoint 函數之外分離 tensors。

參數

• function – 描述在模型或模型的一部分的正向傳播中要運行的內容。它也應該知道如何處理作為 tuple 傳遞的輸入。例如，在 LSTM 中，如果使用者傳遞 (activation, hidden)，則 function 應該正確地使用第一個輸入作為 activation，並將第二個輸入作為 hidden

• preserve_rng_state (bool, 可選) – 省略在每次檢查點期間儲存和恢復 RNG 狀態。請注意，在 torch.compile 下，此標誌不起作用，我們始終會保留 RNG 狀態。預設值：True

• use_reentrant (bool) – 指定是否使用需要可重入自動微分 (reentrant autograd) 的激活檢查點變體。此參數應明確傳遞。在 2.5 版本中，如果未傳遞 use_reentrant，我們將引發異常。如果 use_reentrant=False，checkpoint 將使用不需要可重入自動微分的實作。這允許 checkpoint 支援額外的功能，例如與 torch.autograd.grad 搭配使用，並支援傳遞到檢查點函數的關鍵字引數。

• context_fn (Callable, optional) – 一個可呼叫物件，返回一個包含兩個上下文管理器 (context manager) 的元組。該函數及其重新計算將分別在第一個和第二個上下文管理器下運行。此引數僅在 use_reentrant=False 時支援。

• determinism_check (str, optional) – 一個字串，用於指定要執行的決定性檢查 (determinism check)。預設情況下，它被設定為 "default"，將重新計算的張量的形狀、dtypes 和裝置與已保存的張量進行比較。要關閉此檢查，請指定 "none"。目前僅支援這兩個值。如果您希望看到更多決定性檢查，請提出 issue。此引數僅在 use_reentrant=False 時支援，如果 use_reentrant=True，則始終禁用決定性檢查。

• debug (bool, optional) – 如果為 True，錯誤訊息還將包含在原始正向計算和重新計算期間運行的運算符的追蹤資訊。此引數僅在 use_reentrant=False 時支援。

• args – 包含 function 的輸入的元組

Returns

在 *args 上運行 function 的輸出

torch.utils.checkpoint.checkpoint_sequential(functions, segments, input, use_reentrant=None, **kwargs)

對序列模型進行檢查點處理，以節省記憶體。

序列模型按順序（循序地）執行模組/函數的清單。因此，我們可以將這樣的模型劃分為多個片段 (segment)，並對每個片段進行檢查點處理。除了最後一個片段之外的所有片段都不會儲存中間激活值。每個檢查點片段的輸入將被保存，以便在反向傳遞中重新運行該片段。

警告

應該明確傳遞 use_reentrant 參數。在 2.4 版本中，如果未傳遞 use_reentrant，我們將引發異常。如果您正在使用 use_reentrant=True` 變體，請參閱 :func:`~torch.utils.checkpoint.checkpoint` 以了解此變體的重要考量和限制。建議您使用``use_reentrant=False。

參數

• functions – 一個 torch.nn.Sequential 或模組或函數（組成模型）的清單，以循序方式運行。

• segments – 模型中要建立的區塊數量

• input – 輸入到 functions 的 Tensor

• preserve_rng_state (bool, optional) – 省略在每個檢查點期間儲存和恢復 RNG 狀態。預設值：True

• use_reentrant (bool) – 指定是否使用需要可重入自動微分 (reentrant autograd) 的激活檢查點變體。此參數應明確傳遞。在 2.5 版本中，如果未傳遞 use_reentrant，我們將引發異常。如果 use_reentrant=False，checkpoint 將使用不需要可重入自動微分的實作。這允許 checkpoint 支援額外的功能，例如與 torch.autograd.grad 搭配使用，並支援傳遞到檢查點函數的關鍵字引數。

Returns

在 *inputs 上循序運行 functions 的輸出

Example

>>> model = nn.Sequential(...)
>>> input_var = checkpoint_sequential(model, chunks, input_var)

torch.utils.checkpoint.set_checkpoint_debug_enabled(enabled)

上下文管理器，用於設定檢查點在運行時是否應印出額外的偵錯資訊。有關更多資訊，請參閱 checkpoint() 的 debug 標誌。請注意，設定後，此上下文管理器將覆蓋傳遞給檢查點的 debug 值。要延遲到本地設定，請將 None 傳遞給此上下文。

參數

enabled (bool) – 檢查點是否應印出偵錯資訊。預設值為 ‘None’。

class torch.utils.checkpoint.CheckpointPolicy(value)

用於指定反向傳播期間檢查點處理策略的枚舉。

支援以下策略

• {MUST,PREFER}_SAVE：表示此操作的輸出將在前向傳遞期間儲存，並且在反向傳遞期間不會重新計算。

• {MUST,PREFER}_RECOMPUTE：表示此操作的輸出將在前向傳遞期間不儲存，並且在反向傳遞期間將會重新計算。

使用 MUST_* 優於 PREFER_*，以表明該策略不應被其他子系統（如 torch.compile）覆蓋。

注意

始終返回 PREFER_RECOMPUTE 的策略函數等同於一般的檢查點 (vanilla checkpointing)。

每個操作都返回 PREFER_SAVE 的策略函數不等同於不使用檢查點。使用這種策略會儲存額外的張量，而不僅限於實際梯度計算所需的張量。

class torch.utils.checkpoint.SelectiveCheckpointContext(*, is_recompute)

在選擇性檢查點期間傳遞到策略函數的上下文。

此類別用於在選擇性檢查點期間將相關元數據傳遞到策略函數。 元數據包括策略函數的目前調用是否在重新計算期間。

Example

>>>
>>> def policy_fn(ctx, op, *args, **kwargs):
>>>    print(ctx.is_recompute)
>>>
>>> context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn)
>>>
>>> out = torch.utils.checkpoint.checkpoint(
>>>     fn, x, y,
>>>     use_reentrant=False,
>>>     context_fn=context_fn,
>>> )

torch.utils.checkpoint.create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False)

輔助函數，用於避免在激活檢查點期間重新計算某些操作。

將此函數與 torch.utils.checkpoint.checkpoint 搭配使用，以控制在反向傳遞期間重新計算哪些操作。

參數

• policy_fn_or_list (Callable or List) –

  • 如果提供了策略函數，它應該接受一個 SelectiveCheckpointContext、OpOverload、操作的 args 和 kwargs，並返回一個 CheckpointPolicy 枚舉值，指示是否應重新計算該操作的執行。

  • 如果提供了一個操作列表，則相當於一個策略，該策略對指定的操作返回 CheckpointPolicy.MUST_SAVE，對所有其他操作返回 CheckpointPolicy.PREFER_RECOMPUTE。

• allow_cache_entry_mutation (bool, optional) – 預設情況下，如果選擇性激活檢查點所快取的任何張量被改變，為了確保正確性，會引發錯誤。如果設定為 True，則會停用此檢查。

Returns

一個包含兩個上下文管理器的 tuple。

Example

>>> import functools
>>>
>>> x = torch.rand(10, 10, requires_grad=True)
>>> y = torch.rand(10, 10, requires_grad=True)
>>>
>>> ops_to_save = [
>>>    torch.ops.aten.mm.default,
>>> ]
>>>
>>> def policy_fn(ctx, op, *args, **kwargs):
>>>    if op in ops_to_save:
>>>        return CheckpointPolicy.MUST_SAVE
>>>    else:
>>>        return CheckpointPolicy.PREFER_RECOMPUTE
>>>
>>> context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn)
>>>
>>> # or equivalently
>>> context_fn = functools.partial(create_selective_checkpoint_contexts, ops_to_save)
>>>
>>> def fn(x, y):
>>>     return torch.sigmoid(torch.matmul(torch.matmul(x, y), y)) * y
>>>
>>> out = torch.utils.checkpoint.checkpoint(
>>>     fn, x, y,
>>>     use_reentrant=False,
>>>     context_fn=context_fn,
>>> )