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