经过测试,通过sess.close释放资源并不能释放显存,导致重新建立一个session分配资源时失败。
stackoverflow上提到了两种方法:
1.利用numba库(需要安装numba库)
背后原理:
Tensorflow只是为GPU分配显存,而CUDA负责管理GPU内存。
如果在用sess.close()清除了所有图形后,CUDA以某种方式拒绝释放GPU显存,那么您可以使用cuda库直接控制CUDA以清除GPU显存。
from numba import cudadef train():...train()
cuda.select_device(0) # 选择GPU设备
cuda.close() # 释放GPU资源
缺点:
以这种方式清除GPU显存会存在问题,再次开启session时会失败,原因貌似是cuda.close破坏了tensorflow开启session的环境。
from numba import cudadef train():...train()
cuda.select_device(0) # 选择GPU设备
cuda.close() # 释放GPU资源
train() # 失败(抛出异常)
2.利用multiprocessing的Process
背后原理:
进程结束后会自动释放和进程相关的资源,这就好比模型训练完成退出程序后会自动释放显存。
from multiprocessing import Processdef train():...if __name__'__main__': # 在windows必须在这句话下面开启多进程p = Process(target=train)p.start()p.join() # 进程结束后,GPU显存会自动释放p = Process(target=train) # 重新训练p.start()p.join()
优点:
相比第一种方法,可以在释放GPU显存后重新开启session。
参考资料:
stackoverflow