深度学习模型函数执行完毕后显存不释放问题解决方案

使用多线程的方法来解决

如下定义model_predict()函数为调用深度学习模型预测结果函数,返回接口列表,使用多线程的方法可以使得调用函数后释放显卡,以便在同张显卡上调用第二个深度学习模型。
具体代码如下

import multiprocessing

# return_dict用于函数返回值
def model_predict(X_test, return_dict):
	...
	...
	res = predict(X_test)
	return_dict["res"] = res
	
def main():
	# 获取多线程中函数的返回值
    manager = multiprocessing.Manager()
    return_dict = manager.dict()
	#X_test 为模型输入
	
	p = multiprocessing.Process(target=model_predict, args=(X_test, return_dict))
	p.start()
	p.join()
	res = return_dict["res"]

你可能感兴趣的:(人工智能,python,显卡释放,多线程,深度学习)