TensorFlow函数使用完成释放显存问题
使用Tensorflow跑程序,3G显存总是会占满,之后使用Pytorch就会显存不足。Tensorflow部分使用完成,显存不会自动释放,直到程序结束才会释放。我们可以借助多进程来释放内存。下面是我尝试各种办法后,得到的目前最好方法。
其余办法不是太慢,就是无法输入多个参数,或者得不到返回值
import multiprocessing
def extract_regional_features(rgb_list, model_path):
reg_feat_list = []
model = get_model('reg_model')(model_path)
for _, val in enumerate(rgb_list):
reg_feat = model.run_test_data(val)
reg_feat_list.append(reg_feat)
model.close()
return reg_feat_list
pool = multiprocessing.Pool(processes=1)
result = pool.apply_async(extract_regional_features, args=(rgb_list, reg_model_path))
pool.close()
pool.join()
print(result)
A = result.get()