【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,那么程序就会默认把整块卡的所有显存占满。