tensorflow-gpu使用

2019-03-24  本文已影响0人  脏脏包盛

python代码中设置使用的GPU

如果要在python代码中设置使用的GPU(如使用pycharm进行调试时),可以使用下面的代码

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

设置tensorflow使用的显存大小

定量设置显存

默认tensorflow是使用GPU尽可能多的显存。可以通过下面的方式,来设置使用的GPU显存:

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

上面分配给tensorflow的GPU显存大小为:GPU实际显存*0.7。

按需设置显存

上面的只能设置固定的大小。如果想按需分配,可以使用allow_growth参数

gpu_options = tf.GPUOptions(allow_growth=True)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) 
上一篇 下一篇

猜你喜欢

热点阅读