跑模型——多次推理后显存累计的问题

前言

做项目的时候，发现随着多张图片进行并行计算的时候，占用的显存没有随着程序的结束而自动消除掉申请的显存

Pytorch显存分析(不要用nvidia-smi) / 显存释放机制

      因为Pytorch的机制是使用缓存分配器来管理缓存分配的(因为这样速度快), 但是在缓存分配器的机制下, 一个Tensor就算被释放了，进程也不会把空闲出来的显存还给GPU，而是等待下一个Tensor来填入这一片被释放的空间(即只要一个Tensor对象在后续不会再被使用，那么PyTorch就会自动回收该Tensor所占用的显存，并以缓冲区的形式继续占用显存，所以在nvidia-smi/gpustat中看到的显存并没有减少)，而且nvidia-smi显示的已经分配到的显存(显存缓冲区)和context开销之和

解决方法

在网上查找资料，在PyTorch中，GPU内存的释放并不总是自动完成的。

虽然CUDA有垃圾回收机制，可以在数据队列达到一定阈值时触发清理失活内存，但这个过程不是立即执行的。因此，为了确保GPU内存得到有效管理，通常需要手动干预。具体如下：

1、使用torch.cuda.empty_cache()

这是一个简单而有效的方法，用于清理PyTorch模型训练后GPU的缓存。它会释放当前进程持有的所有未使用的缓存内存，有助于避免内存泄漏。

2、调整批量大小（batch size）

减小批量大小可以减少每次推理时的内存使用，从而降低GPU内存的压力，训练的时候，内存不足可以这么做哟

3、手动释放Tensor对象

在使用完模型后，及时删除不再使用的Tensor和模型，以便它们的内存可以被回收

最快最方便的就是调用torch.cuda.empty_cache()函数来释放GPU内存。这个操作会清空PyTorch未使用的显存缓存，帮助减少GPU内存的占用，放在推理代码使用完GPU的代码之后就行了

class your_model():
    def __init__(self):
        self.net = model(model_path)

    def detect_image(image_data):
        with torch.no_grad():
            images = torch.from_numpy(image_data)
            images = images.cuda()
            # ---------------------------------------------------#
            #   图片传入网络进行预测
            # ---------------------------------------------------#
            pr = self.net(images)[0]

            torch.cuda.empty_cache()
            return pr

注：在预测的时候，我们不需要进行梯度的计算，要使用torch.no_grad()上下文管理器来避免梯度的存储哦~

我觉得在执行完GPU的计算之后释放GPU内存主要看需求吧，毕竟会使程序速度变慢一点~

欢迎大家点赞或收藏~

点赞或收藏可以鼓励作者加快更新哟~

参考链接：

Pytorch显存机制与显存占用(一)