【tensorflow】设置显存开销

问题

一般大家在跑tf时,单个程序往往会占满整块GPU的所有显存。
但是实际上,程序很可能并不需要那么大的显存空间。

改进方案

通过 tf.ConfigProto().gpu_options.allow_growth=True来告诉程序,只需占用实际所需的显存即可:

# ---------------- session ----------------
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)

如果这里把 config.gpu_options.allow_growth设置为False,那么程序就会默认把整块卡的所有显存占满。

发布了628 篇原创文章 · 获赞 863 · 访问量 185万+

猜你喜欢

转载自blog.csdn.net/JNingWei/article/details/104848722