4.6_use-gpu.md

GPU计算

到目前为止，我们一直在使用CPU计算。对复杂的神经网络和大规模的数据来说，使用CPU来计算可能不够高效。在本节中，我们将介绍如何使用单块NVIDIA GPU来计算。首先，需要确保已经安装好了至少一块NVIDIA GPU。然后，下载CUDA并按照提示设置好相应的路径

注意：需要tensorflow-gpu

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')]

check available device

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"
]

specify device

使用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)

注：本节与原书有很大不同，原书传送门