GPU的使用说明

查看显卡信息

!nvidia-smi

计算设备

import tensorflow as tf

tf.device('/CPU:0'), tf.device('/GPU:0'), tf.device('/GPU:1')

查询GPU数量

len(tf.config.experimental.list_physical_devices('GPU'))

张量与GPU
我们可以查询张量所在的设备。默认情况下，张量是在CPU上创建的。

x = tf.constant([1, 2, 3])
x.device

存储在GPU上
我们需要确保不创建超过GPU显存限制的数据。

with try_gpu():
    X = tf.ones((2, 3))

假设你至少有两个GPU，下面的代码将在第二个GPU上创建一个随机张量。

with try_gpu(1):
    Y = tf.random.uniform((2, 3))

复制

with try_gpu(1):
    Z = X
print(X)
print(Z)

现在数据在同一个GPU上（Z和Y都在），我们可以将它们相加。

假设变量Z已经存在于第二个GPU上。如果我们仍然在同一个设备作用域下调用Z2 = Z会发生什么？它将返回Z，而不会复制并分配新内存。

with try_gpu(1):
    Z2 = Z
    
Z2 is Z
#True

神经网络与GPU
类似地，神经网络模型可以指定设备。下面的代码将模型参数放在GPU上。

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    net = tf.keras.models.Sequential([
        tf.keras.layers.Dense(1)])

当输入为GPU上的张量时，模型将在同一GPU上计算结果。

net(X)

让我们确认模型参数存储在同一个GPU上。

net.layers[0].weights[0].device, net.layers[0].weights[1].device

所有的数据和参数都在同一个设备上，我们就可以有效地学习模型

小结

我们可以指定用于存储和计算的设备，例如CPU或GPU。默认情况下，数据在主内存中创建，然后使用CPU进行计算。
深度学习框架要求计算的所有输入数据都在同一设备上，无论是CPU还是GPU。
不经意地移动数据可能会显著降低性能。一个典型的错误如下：计算GPU上每个小批量的损失，并在命令行中将其报告给用户（或将其记录在NumPy ndarray中）时，将触发全局解释器锁，从而使所有GPU阻塞。最好是为GPU内部的日志分配内存，并且只移动较大的日志。