指定tensorflow运行的GPU或CPU设备

程序员文章站 2022-04-03 09:58:44

如果 TensorFlow 指令中兼有 CPU 和 GPU 实现，当该指令分配到设备时，GPU 设备有优先权。如果你的系统里有多个 GPU, 那么 ID 最小的 GPU 会默认使用。当我们要指定tensorflow运行的GPU或CPU设备时，可以使用tf.device()命令首先查看可用运算设备(CPU,GPU)from tensorflow.python.client import device_libprint(device_lib.list_local_devices())得.....

如果 TensorFlow 指令中兼有 CPU 和 GPU 实现，当该指令分配到设备时，GPU 设备有优先权。
如果你的系统里有多个 GPU, 那么 ID 最小的 GPU 会默认使用。

当我们要指定tensorflow运行的GPU或CPU设备时，可以使用tf.device()命令

首先查看可用运算设备(CPU,GPU)

from tensorflow.python.client import device_lib
print(device_lib.list_local_devices())

得到类似以下的输出结果

[name: "/device:CPU:0"
device_type: "CPU"
memory_limit: 268435456
locality {
}
incarnation: 13177083330855175469
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 10968950375
locality {
  bus_id: 1
}
incarnation: 6161624703599064583
physical_device_desc: "device: 0, name: GeForce GTX 1080 Ti, pci bus id: 0000:00:08.0, compute capability: 6.1"
]

name即是对应设备名称，一般来说（以各自实际情况为准，每个人的情况可能不同）

"/cpu:0"：机器的 CPU。
"/device:GPU:0"：机器的 GPU（如果有一个）。
"/device:GPU:1"：机器的第二个 GPU（以此类推）。

使用命令tf.device()进行指定

import os
os.environ["CUDA_VISIBLE_DEVICES"] = '/device:GPU:0'

本文地址：https://blog.csdn.net/tsundere_x/article/details/107667896