本文简要介绍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
参数:
相关用法
- Python PyTorch saved_tensors_hooks用法及代码示例
- Python PyTorch save用法及代码示例
- Python PyTorch sqrt用法及代码示例
- Python PyTorch skippable用法及代码示例
- Python PyTorch squeeze用法及代码示例
- Python PyTorch square用法及代码示例
- Python PyTorch scatter_object_list用法及代码示例
- Python PyTorch skip_init用法及代码示例
- Python PyTorch simple_space_split用法及代码示例
- Python PyTorch sum用法及代码示例
- Python PyTorch sub用法及代码示例
- Python PyTorch sparse_csr_tensor用法及代码示例
- Python PyTorch sentencepiece_numericalizer用法及代码示例
- Python PyTorch symeig用法及代码示例
- Python PyTorch sinh用法及代码示例
- Python PyTorch sinc用法及代码示例
- Python PyTorch std_mean用法及代码示例
- Python PyTorch spectral_norm用法及代码示例
- Python PyTorch slogdet用法及代码示例
- Python PyTorch symbolic_trace用法及代码示例
- Python PyTorch shutdown用法及代码示例
- Python PyTorch sgn用法及代码示例
- Python PyTorch set_flush_denormal用法及代码示例
- Python PyTorch set_default_dtype用法及代码示例
- Python PyTorch signbit用法及代码示例
注:本文由纯净天空筛选整理自pytorch.org大神的英文原创作品 torch.autograd.graph.save_on_cpu。非经特殊声明,原始代码版权归原作者所有,本译文未经允许或授权,请勿转载或复制。