Python PyTorch save_on_cpu用法及代碼示例

本文簡要介紹python語言中 torch.autograd.graph.save_on_cpu 的用法。

用法:

class torch.autograd.graph.save_on_cpu(pin_memory=False)

參數：

pin_memory(bool) -如果 True 張量將在打包期間保存到 CPU 固定內存，並在解包期間異步複製到 GPU。默認為 False 。另請參閱使用固定內存緩衝區。

Context-manager 由前向傳遞保存的張量將存儲在 cpu 上，然後向後檢索。

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

使用此 context-manager 來交換 GPU 內存使用的計算(例如，當您的模型在訓練期間不適合 GPU 內存時)。

例子：

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