释放pytorch占用的gpu显存_使用GPU训练大神经网络之显存控制策略与提升

最新推荐文章于 2024-08-29 03:32:27 发布

竹谭

最新推荐文章于 2024-08-29 03:32:27 发布

阅读量694

点赞数

文章标签：释放pytorch占用的gpu显存

本文链接：https://blog.csdn.net/weixin_42357331/article/details/112282091

版权

本文介绍了PyTorch中优化GPU显存管理的策略，包括资源回收、数据迁移和重计算。论文SuperNeurons探讨了如何通过Live Analysis、Offloading & Prefetching和Recomputation提高训练深度学习模型时的显存利用率，以支持更大模型或批处理大小。同时，论文中对比分析了Caffe、TensorFlow、MXNet和Torch的显存控制策略，提出了动态的显存控制策略的重要性。

摘要由CSDN通过智能技术生成

本文是“[PPoPP18] SuperNeurons:Dynamic GPU Memory Management for Training Deep Neural Networks”的论文阅读笔记，论文、talk视频PPT等详见作者主页：https://linnanwang.github.io/
参考资料：https://linnanwang.github.io/
How to Optimize Data Transfers in CUDA C/C++
How to Overlap Data Transfers in CUDA C/C++
Memory Allocation Overhead
cudamalloc slow
vDNN: Virtualized Deep Neural Networks for Scalable, Memory-Efficient Neural Network Design
Training Deep Nets with Sublinear Memory Cost
Arbitrary 2D convolution
Fast Algorithms for Convolutional Neural Networks

GPU显存控制策略与提升

该论文对时下各个框架中显存控制的思想做了汇总，并做了一些技术上的提升。整体上分为三个部分：
1. Live Analysis：即资源回收，回收无用tensor的显存。
2. offloading&prefetching：将装不下的tensor暂存到内存等其他设备中。
3. Recomputation：对于占用显存多、计算消耗小的tensor，不进行存储，需要使用时重新计算。

1. Live Analysis

其思想主要是资源回收，对于后续操作不依赖的tensor进行回收，释放相应的显存空间。
关于如何分析依赖关系，笔者猜测应该有许多基于计算图的算法。论文中给出了一个简单的O(N2)O(N2)算法（NN是网络层数）：维护一个需要保留的tensor集合，每层计算结束后加入当前层输出，并判断集合中的tensor是否被后续计算所依赖，如不被依赖则删除。
性能提升：简单的新增和删除tensor涉及到许多的cudaMalloc和cudaFree操作，会占用较多时间。可以一次性申请足够多的空间，然后程序自身维护空间内的赋值，标记tensor的活跃度等操作。详情请参考：Memory Allocation Overhead ，How to Optimize Data Transfers in CUDA C/C++。

2. Offloading & Prefetching

其思想主要是将GPU中装不下的tensor暂存到内存等其他设备中。论文中的Unified Tensor Pool (UTP)说的是其可以应用于各种与Local GPU可以交互的设备之间，但实际实现时该runtime仅对CPU内存和GPU显存之间做传输。
有效的GPU相关数据传输技巧详见：How to Optimize Data Transfers