到目前为止,我们一直在使用CPU计算。对复杂的神经网络和大规模的数据来说,使用CPU来计算可能不够高效。在本节中,我们将介绍如何使用单块NVIDIA GPU来计算。首先,需要确保已经安装好了至少一块NVIDIA GPU。然后,下载CUDA并按照提示设置好相应的路径
import tensorflow as tf
import numpy as np
print(tf.__version__)
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
cpus = tf.config.experimental.list_physical_devices(device_type='CPU')
print("可用的GPU:",gpus,"\n可用的CPU:", cpus)
2.0.0
可用的GPU: [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
可用的CPU: [PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]
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: 15592483132577835191
, name: "/device:GPU:0"
device_type: "GPU"
memory_limit: 3063309926
locality {
bus_id: 1
links {
}
}
incarnation: 3859243074925251015
physical_device_desc: "device: 0, name: GeForce GTX 1650, pci bus id: 0000:01:00.0, compute capability: 7.5"
]
使用tf.device()来指定特定设备(GPU/CPU)
with tf.device('GPU:0'):
a = tf.constant([1,2,3],dtype=tf.float32)
b = tf.random.uniform((3,))
print(tf.exp(a + b) * 2)
tf.Tensor([12.172885 19.682476 53.18001 ], shape=(3,), dtype=float32)
注:本节与原书有很大不同,原书传送门