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

mac2024-04-02  40

使用多线程的方法来解决

如下定义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"]
最新回复(0)