释放pytorch占用的gpu显存_再次浅谈Pytorch中的显存利用问题(附完善显存跟踪代码)...-CSDN博客

本文链接：https://blog.csdn.net/weixin_36080429/article/details/112042990

本文介绍了如何使用Pytorch-Memory-Utils工具检测和分析训练过程中GPU显存的使用情况。通过实例展示了模型权重、额外显存开销以及临时缓冲值对显存的影响。讨论了PyTorch的显存回收机制，以及`torch.no_grad()`在减少显存占用中的作用。最后提到了异步执行对显存管理的影响，并鼓励读者关注PyTorch的更新以优化显存利用。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

前言

之前在浅谈深度学习:如何计算模型以及中间变量的显存占用大小和如何在Pytorch中精细化利用显存中我们已经谈论过了平时使用中显存的占用来自于哪里，以及如何在Pytorch中更好地使用显存。在这篇文章中，我们借用Pytorch-Memory-Utils这个工具来检测我们在训练过程中关于显存的变化情况，分析出我们如何正确释放多余的显存。

在深度探究前先了解下我们的输出信息，通过Pytorch-Memory-Utils工具，我们在使用显存的代码中间插入检测函数(如何使用见工具github页面和下文部分)，就可以输出类似于下面的信息，At __main__ : line 13 Total Used Memory:696.5 Mb表示在当前行代码时所占用的显存，即在我们的代码中执行到13行的时候所占显存为695.5Mb。At __main__ : line 15 Total Used Memory:1142.0 Mb表示程序执行到15行时所占的显存为1142.0Mb。两条数据之间表示所占显存的tensor变量。

# 12-Sep-18-21:48:45-gpu_mem_track.txt

GPU Memory Track | 12-Sep-18-21:48:45 | Total Used Memory:696.5 Mb

At __main__ : line 13 Total Used Memory:696.5 Mb

+ | 7 * Size:(512, 512, 3, 3) | Memory: 66.060 M |

+ | 1 * Size:(512, 256, 3, 3) | Memory: 4.7185 M |

+ | 1 * Size:(64, 64, 3, 3) | Memory: 0.1474 M |

+ | 1 * Size:(128, 64, 3, 3) | Memory: 0.2949 M |

+ | 1 * Size:(128, 128, 3, 3) | Memory: 0.5898 M |

+ | 8 * Size:(512,) | Memory: 0.0163 M |

+ | 3 * Size:(256, 256, 3, 3) | Memory: 7.0778 M |

+ | 1 * Size:(256, 128, 3, 3) | Memory: 1.1796 M |

+ | 2 * Size:(64,) | Memory: 0.0005 M |