tensorflow释放显存

最新推荐文章于 2026-04-17 04:18:07 发布

原创最新推荐文章于 2026-04-17 04:18:07 发布 · 1.4w 阅读

21 ·

本内容遵循CC 4.0 BY-SA版权协议

标签

#cuda #tensorflow #深度学习

本文探讨了在Tensorflow中释放GPU显存的两种方法：使用numba库的cuda模块直接控制CUDA清理显存，以及利用multiprocessing的Process特性在进程结束后自动释放资源。后者可在释放GPU显存后重新开启session，更为实用。

经过测试，通过sess.close释放资源并不能释放显存，导致重新建立一个session分配资源时失败。

stackoverflow上提到了两种方法：

1.利用numba库（需要安装numba库）

背后原理：
Tensorflow只是为GPU分配显存，而CUDA负责管理GPU内存。
如果在用sess.close()清除了所有图形后，CUDA以某种方式拒绝释放GPU显存，那么您可以使用cuda库直接控制CUDA以清除GPU显存。

from numba import cuda

def train():
    .
    .
    .


train()
cuda.select_device(0) # 选择GPU设备
cuda.close() # 释放GPU资源

缺点：
以这种方式清除GPU显存会存在问题，再次开启session时会失败，原因貌似是cuda.close破坏了tensorflow开启session的环境。

from numba import cuda

def train():
    .
    .
    .


train()
cuda.select_device(0) # 选择GPU设备
cuda.close() # 释放GPU资源
train() # 失败（抛出异常）

2.利用multiprocessing的Process

背后原理：
进程结束后会自动释放和进程相关的资源，这就好比模型训练完成退出程序后会自动释放显存。

from multiprocessing import Process

def train():
    .
    .
    .

if __name__'__main__': # 在windows必须在这句话下面开启多进程
    p = Process(target=train)
    p.start()
    p.join() # 进程结束后，GPU显存会自动释放
    p = Process(target=train) # 重新训练
    p.start()
    p.join()