python_debug系列: TensorFlow模型级联问题(同一个网络,加载不同参数,对不同数据进行预测)

本文讲述了在TensorFlow中遇到的模型级联问题,即如何在不释放显存的情况下,使用同一个网络结构但加载不同参数对不同数据进行预测。作者通过尝试关闭Session、使用numba清空显存以及重置默认计算图,最终解决了问题,实现了模型的级联预测。
摘要由CSDN通过智能技术生成

今天在工作中遇到了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

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值