给我一个TensorFlow 2.x 可以用的测试GPU的代码
时间: 2024-01-24 07:17:40 浏览: 22
以下是一个使用 TensorFlow 2.x 测试 GPU 的代码示例:
```python
import tensorflow as tf
# 检查是否有可用的 GPU
print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
# 创建一个 GPU 设备对象,如果没有可用的 GPU 则会使用 CPU
device_name = tf.test.gpu_device_name()
if device_name != '/device:GPU:0':
device_name = '/device:CPU:0'
# 在设备上运行一些 TensorFlow 操作
with tf.device(device_name):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
# 输出结果
print(c)
```
运行以上代码,如果你的电脑上有可用的 GPU,则会输出类似以下内容:
```
Num GPUs Available: 1
tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32)
```
其中 `Num GPUs Available` 表示可用的 GPU 数量,`tf.Tensor` 表示 TensorFlow 的计算结果。如果你的电脑上没有可用的 GPU,则会使用 CPU 运行,输出类似以下内容:
```
Num GPUs Available: 0
tf.Tensor(
[[22. 28.]
[49. 64.]], shape=(2, 2), dtype=float32)
```