今天在工作中遇到了TensorFlow模型级联问题,主要涉及的问题是显存和计算图的初始化
问题描述:
在训练一个3D-VNet模型时,网络结构是封装在VNet3D类中,在类内定义了init,train, predict, re_train等函数。使用相同数据集将其降采样为不同的分辨率后训练了几个不同的模型,而在实际预测时,需要将测试集分别经过低分辨率模型和高分辨率模型,用于定位和精确重建。
每个模型在单独测试时没有问题,但是将两个模型进行级联时,python便会报错。
在jupyternotebook进行调试时发现在调用完predict函数并完成模型1预测任务后,TensorFlow并不会释放显存:
初步想法是在第一次预测结束后,关闭session,释放显存,然后载入第二个模型,进行预测
- 尝试1
在类中添加关闭session的函数
def close_sess(self):
self.sess.close()
运行后,依然报错。session似乎关闭了,但是依然没有释放显存
- 尝试2
使用numba来调用cuda库直接清空显存
from numba import cuda
cuda.select_device(0) # 选择GPU设备
cuda.c