python

关注公众号 jb51net

关闭
首页 > 脚本专栏 > python > tensorflow 限制显存大小

tensorflow 限制显存大小的实现

作者:安阳小栈-客官歇会吧

今天小编就为大家分享一篇tensorflow 限制显存大小的实现,具有很好的参考价值,希望对大家有所帮助。一起跟随小编过来看看吧

Python在用GPU跑模型的时候最好开多进程,因为很明显这种任务就是计算密集型的。

用进程池好管理,但是tensorflow默认情况会最大占用显存,尽管该任务并不需要这么多,因此我们可以设置显存的按需获取,这样程序就不会死掉了。

1. 按比例预留:

tf_config = tensorflow.ConfigProto() 
tf_config.gpu_options.per_process_gpu_memory_fraction = 0.5 # 分配50% 
session = tensorflow.Session(config=tf_config) 

2. 或者干脆自适应然后自动增长:

tf_config = tensorflow.ConfigProto() 
tf_config.gpu_options.allow_growth = True # 自适应 
session = tensorflow.Session(config=tf_config) 

以上这篇tensorflow 限制显存大小的实现就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持脚本之家。

您可能感兴趣的文章:
阅读全文