【tensorflow】设置显存开销
- 2020 年 3 月 17 日
- 笔记
问题
一般大家在跑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
,那么程序就会默认把整块卡的所有显存占满。