tensorflow 2.0 gpu集可用gpu？

Question

from tensorflow.python.client import device_lib

def get_available_gpus():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos if x.device_type == 'GPU']
get_available_gpus()

'/device: gpu :0'，/device: gpu :1‘当我用tensorflow 2.0 gpu训练和编译cnn时，我想设置可用的gpu 我该怎么做？

os.environ["CUDA_VISIBLE_DEVICES"]="1";  
from keras.backend.tensorflow_backend import set_session
config = tf.compat.v1.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.8
tf.compat.v1.Session(config=config)

有些人认为这段代码是不可靠的

tf.版‘2.0’

Answer 1

据我在你的脚本中所看到的，它看起来已经相当正确了。据我所知，tensorflow 2.x的keras有自己的高级API，称为keras。我也用

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

和

try:
  %tensorflow_version 2.x
except Exception:
  pass
import tensorflow as tf
tf.keras.backend.clear_session()

另外，您可能需要切换到

from tensorflow import keras
from tensorflow.keras.optimizers import Adam
from tensorflow.keras import layers

您可以从这个链接中找到原因。我认为首先，您可以使用预先构建的(例如VGG或ResNet)模型进行测试，以查看您的GPU是否运行。我非常肯定，我的一些定制构建模型并不真正使用GPU。但是当我和另一种型号一起跑的时候，它就起作用了。这有点奇怪，问题就在我假设的模型上。

Answer 2

使用with tf.device():

因此，您现在有两个具有名称的GPU：

['/device:GPU:0', '/device:GPU:1'] 

然后在您的代码中，例如您的主要工作：

def main():
    # do somthing here...

主要()：

if __name__ == "__main__":
    with tf.device('/device:GPU:0'): # select device here...
        main()