关于TensorFlow的GPU设置

摘自:https://blog.csdn.net/byron123456sfsfsfa/article/details/79811286

1.  在使用GPU版的TensorFlow跑程序的时候,如果不特殊写代码注明,程序默认是占用所有主机上的GPU,但计算过程中只会用其中一块。也就是你看着所有GPU都被占用了,以为是在GPU并行计算,但实际上只有其中一块在运行;另外的所有显卡都闲着,但其显存都被占用了,所以别人也用不了。不过这种情况通过在程序之前加三行代码就可以解决:

1 import os
2 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
3 # 保证程序中的GPU序号是和硬件中的序号是相同的
4 
5 os.environ['CUDA_VISIBLE_DEVICES'] = "0,1"
6 # 屏蔽掉系统中除 gpu0 和 gpu1 (gpu序号要根据实际情况来定)之外所有的GPU设备
7 # 即设置仅gpu0,gpu1可见
8 # os.environ['CUDA_VISIBLE_DEVICES'] = "" 即不使用gpu

2. 避免TensorFlow程序会吃掉所用显卡的所有显存,设置程序需要多少显存就用多少

1 # TensorFlow程序会吃掉所用显卡的所有显存
2 # 如果想让程序需要多少显存就用多少
3 # 创建 session 的时候加一项设置
4 config = tf.ConfigProto(log_device_placement=FLAGS.log_device_placement)
5 config.gpu_options.allow_growth = True
6 session = tf.Session(config=config)

猜你喜欢

转载自www.cnblogs.com/shiyublog/p/10085140.html