解决tensorflow会话结束后,GPU资源不释放

在使用tensorflow对医疗影像x-ray胸片进行在线病理诊断时发现,在进行一次图片处理之后GPU占用率一直被服务进程100%占用,资源不会自动释放。tensorflow issue上查找问题发现,这个问题原本就是tf设计上的一个bug,问题还没有被正确修复。https://github.com/tensorflow/tensorflow/issues/1727

问题解决只能每次去杀掉进程来释放资源。而我们是提供的webServer在线提供tf判断。

由于我们的webserver是python做的。基本思路就是再调用TF方法处理时启一个子进程进行方法处理,处理完成后进程结束释放资源。

@app.route('/occlusion/', methods=['POST'])
def getOcclusion(slug):
 filename = os.path.realpath('.') + '/upload/files/'+slug
 return_dict = main(filename,True)
 testprobability=return_dict['probability']
 result = jsonify({'index': np.argmax(testprobability),'value':str(np.max(testprobability)),"hotImageFlag":hotImageFlag})
 return   result

def run_graph(*****, return_dict):
 config = tf.ConfigProto(allow_soft_placement=True,device_count={'gpu':0}) 
 config.gpu_options.per_process_gpu_memory_fraction=0.05
 with tf.Session(config=config) as sess:
 occlusion_array_string = StringIO()

 probability = sess.run(softmax_tensor[:, :],{input_layer_name: occlusion_array_string.getvalue()})
 occlusion_array_string.close()
 return_dict['probability'] = probability

def main(filename,noHotImage):
 #   if argv[1:]:
 #     raise ValueError("Unused Command Line Args: %s" % argv[1:])
 # load graph, which is stored in the default session
 load_graph('x_ray/retrained_graph_pneumonia.pb')

 occlusion_array, image_data = load_image(filename)
 # load labels
 labels = load_labels('x_ray/output_labels_pneumonia.txt')
 # 此处multiprocessing的Manager接收run_graph执行后的结果
 manager = Manager()
 return_dict = manager.dict()

 # 此处multiprocessing的Process使用子进程执行run_graph, run_graph中是tf的操作 
 p = Process(target=run_graph, args=(****,return_dict))
 p.start()
 #子进程处理完后结束进程,这时GPU 资源会释放
 p.join()
 print("return_dict:" + str(return_dict))
 return return_dict

以上为我工程中相关的代码段,不完成提供思路,你只需要重点关注main方法中Process的应用。
上面代码中还需要注意的是,由于我需要将结果返回给前端。这里需要接收子进程中的执行结果,multiprocessing的Manager接收run_graph执行后的结果。并返回给服务接口。

当然还有一种比较低级的方式就是通过以下方式设定固定的gpu占用率。

config.gpu_options.per_process_gpu_memory_fraction=0.05

数风流人物,还看今朝。
更多文章请关注我的博客:https://harveyyeung.github.io

你可能感兴趣的:(解决tensorflow会话结束后,GPU资源不释放)