我的tensorflow版本为2.3.0,我希望测试tensorflow是否使用GPU计算,给出详细代码
时间: 2023-06-14 16:04:12 浏览: 86
tensorflow-gpu测试代码
5星 · 资源好评率100%
你可以使用以下代码来测试TensorFlow是否使用GPU计算:
```python
import tensorflow as tf
# 列出可用的物理设备
physical_devices = tf.config.list_physical_devices()
print("可用的物理设备:", physical_devices)
# 列出可用的GPU设备
gpu_devices = tf.config.list_physical_devices('GPU')
print("可用的GPU设备:", gpu_devices)
# 设置Tensorflow仅使用GPU 0来计算
if len(gpu_devices) > 0:
tf.config.experimental.set_visible_devices(gpu_devices[0], 'GPU')
tf.config.experimental.set_memory_growth(gpu_devices[0], True)
# 测试Tensorflow是否使用GPU计算
with tf.device('/CPU:0'):
a = tf.random.normal([1000, 1000])
b = tf.random.normal([1000, 1000])
c = tf.matmul(a, b)
print(c)
```
这段代码首先列出可用的物理设备和GPU设备,然后设置TensorFlow仅使用第一个GPU来计算,并在CPU上执行一个矩阵乘法操作以测试是否使用GPU计算。 如果一切正常,你应该在输出中看到TensorFlow使用GPU来计算。
阅读全文