tensorflow 使用GPU
时间: 2023-07-05 11:33:53 浏览: 80
在 TensorFlow 中使用 GPU 有两种方式:
1. 自动选择可用的 GPU:如果你的计算机上装有一个或多个 NVIDIA GPU,TensorFlow 会自动选择可用的 GPU。你可以使用以下代码检查 TensorFlow 是否使用了 GPU:
```
import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
```
如果输出结果大于 0,则表示 TensorFlow 正在使用 GPU。
2. 手动选择 GPU:如果你有多个 GPU,或者想指定 TensorFlow 使用哪个 GPU,可以使用以下代码:
```
import tensorflow as tf
# 创建一个 TensorFlow 配置对象
config = tf.compat.v1.ConfigProto()
# 指定使用哪个 GPU,如果有多个 GPU,可以指定多个
config.gpu_options.visible_device_list = "0"
# 创建一个 TensorFlow 会话并指定使用配置对象
sess = tf.compat.v1.Session(config=config)
# 在会话中运行 TensorFlow 代码
...
```
在上面的代码中,我们指定使用第一个 GPU(索引为 0)。如果你有多个 GPU,可以指定多个,例如 `config.gpu_options.visible_device_list = "0,1"`。
阅读全文