【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，那麼程式就會默認把整塊卡的所有顯示記憶體佔滿。