72、使用 GPU 加速 TensorFlow 计算

使用 GPU 加速 TensorFlow 计算

1. 检查 TensorFlow 是否识别 GPU

要确认 TensorFlow 能够识别你的 GPU，可以运行以下命令，确保结果不为空：

>>> physical_gpus = tf.config.list_physical_devices("GPU")
>>> physical_gpus
[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

2. 管理 GPU 内存

默认情况下，TensorFlow 在首次运行计算时会自动占用所有可用 GPU 的几乎全部内存，这是为了减少 GPU 内存碎片化。但这意味着如果尝试启动第二个需要 GPU 的 TensorFlow 程序，可能会很快耗尽内存。不过这种情况并不常见，通常一台机器上只会运行一个 TensorFlow 程序，如训练脚本、TF Serving 节点或 Jupyter 笔记本。

如果因某些原因需要运行多个程序（例如在同一台机器上并行训练两个不同的模型），则需要更均匀地在这些进程之间分配 GPU 内存。以下是几种管理 GPU 内存的方法：
- 为每个进程分配特定的 GPU ：
- 如果你有多张 GPU 卡，可以通过设置 CUDA_VISIBLE_DEVICES 环境变量，让每个进程只看到相应的 GPU 卡。同时，将 CUDA_DEVICE_ORDER 环境变量设置为 PCI_BUS_ID ，以确保每个 ID 始终指向同一张 GPU 卡。
- 例如，如果你有四张 GPU 卡，可以在两个单独的终端窗口中执行以下命令来启动两个程序，每个程序分配两张 GPU：

$ CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0,1 python3 program_1.py
# 在另一个终端中
$ CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=3,2 python3 program_2.py

- 程序 1 在 TensorFlow 中只会看到 GPU 卡 0 和 1，分别命名为 `/gpu:0` 和 `/gpu:1`；程序 2 只会看到 GPU 卡 2 和 3，分别命名为 `/gpu:1` 和 `/gpu:0`。
- 当然，你也可以在 Python 中通过设置 `os.environ["CUDA_DEVICE_ORDER"]` 和 `os.environ["CUDA_VISIBLE_DEVICES"]` 来定义这些环境变量，但必须在使用 TensorFlow 之前进行设置。

• 限制 TensorFlow 占用的 GPU 内存 ：
  • 要让 TensorFlow 只占用每个 GPU 上特定数量的内存，必须在导入 TensorFlow 后立即进行设置。例如，要让 TensorFlow 每个 GPU 只占用 2 GiB 的内存，可以为每个物理 GPU 设备创建一个逻辑 GPU 设备，并将其内存限制设置为 2 GiB（即 2,048 MiB）：

for gpu in physical_gpus:
    tf.config.set_logical_device_configuration(
        gpu,
        [tf.config.LogicalDeviceConfiguration(memory_limit=2048)]
    )

- 假设你有四个 GPU，每个 GPU 至少有 4 GiB 的内存，那么两个这样的程序可以并行运行，每个程序使用所有四个 GPU 卡。运行 `nvidia-smi` 命令时，会看到每个进程在每张卡上占用 2 GiB 的内存。

• 按需分配 GPU 内存 ：
  • 可以让 TensorFlow 仅在需要时占用内存，同样需要在导入 TensorFlow 后立即进行设置：

for gpu in physical_gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

- 另一种方法是将 `TF_FORCE_GPU_ALLOW_GROWTH` 环境变量设置为 `true`。使用此选项时，TensorFlow 一旦占用内存就不会释放（同样是为了避免内存碎片化），除非程序结束。这种方法可能难以保证确定性的行为（例如，一个程序可能因为另一个程序的内存使用量过高而崩溃），因此在生产环境中，可能更倾向于使用前面的选项。但在某些情况下，它非常有用，例如在一台机器上运行多个使用 TensorFlow 的 Jupyter 笔记本时。在 Colab 运行时中，`TF_FORCE_GPU_ALLOW_GROWTH` 环境变量默认设置为 `true`。

• 将一个 GPU 拆分为多个逻辑设备 ：
  • 在某些情况下，你可能想将一个 GPU 拆分为两个或更多逻辑设备。例如，如果你只有一个物理 GPU（如在 Colab 运行时中），但想测试多 GPU 算法，可以使用以下代码将 GPU #0 拆分为两个逻辑设备，每个设备有 2 GiB 的内存：

tf.config.set_logical_device_configuration(
    physical_gpus[0],
    [tf.config.LogicalDeviceConfiguration(memory_limit=2048),
     tf.config.LogicalDeviceConfiguration(memory_limit=2048)]
)

- 这两个逻辑设备分别称为 `/gpu:0` 和 `/gpu:1`，可以像使用两个普通 GPU 一样使用它们。可以使用以下代码列出所有逻辑设备：

>>> logical_gpus = tf.config.list_logical_devices("GPU")
>>> logical_gpus
[LogicalDevice(name='/device:GPU:0', device_type='GPU'),
 LogicalDevice(name='/device:GPU:1', device_type='GPU')]

3. 在设备上放置操作和变量

Keras 和 tf.data 通常能很好地将操作和变量放置在合适的设备上，但如果你需要更多控制，也可以手动将操作和变量放置在每个设备上：
- 一般原则 ：
- 通常应将数据预处理操作放在 CPU 上，将神经网络操作放在 GPU 上。
- GPU 的通信带宽通常较为有限，因此应避免不必要的数据进出 GPU 传输。
- 增加 CPU 内存简单且成本较低，通常有足够的内存；而 GPU 内存是固定在 GPU 中的，是一种昂贵且有限的资源，因此如果一个变量在接下来的几个训练步骤中不需要，可能应该将其放在 CPU 上（例如，数据集通常应放在 CPU 上）。
- 默认设备分配 ：
- 默认情况下，所有变量和操作都会被放置在第一个 GPU（命名为 /gpu:0 ）上，但没有 GPU 内核的变量和操作会被放在 CPU（始终命名为 /cpu:0 ）上。可以通过张量或变量的 device 属性查看其所在的设备：

>>> a = tf.Variable([1., 2., 3.])  # float32 变量会被放在 GPU 上
>>> a.device
'/job:localhost/replica:0/task:0/device:GPU:0'
>>> b = tf.Variable([1, 2, 3])  # int32 变量会被放在 CPU 上
>>> b.device
'/job:localhost/replica:0/task:0/device:CPU:0'

• 手动指定设备 ：
  • 如果你想将操作放置在非默认设备上，可以使用 tf.device() 上下文：

>>> with tf.device("/cpu:0"):
...     c = tf.Variable([1., 2., 3.])
...
>>> c.device
'/job:localhost/replica:0/task:0/device:CPU:0'

- 即使你的机器有多个 CPU 核心，CPU 也始终被视为一个设备（`/cpu:0`）。如果操作有多线程内核，放置在 CPU 上的操作可能会在多个核心上并行运行。
- 如果你明确尝试将操作或变量放置在不存在或没有内核的设备上，TensorFlow 会默默地回退到默认设备。但如果你希望抛出异常，可以运行 `tf.config.set_soft_device_placement(False)`。

4. 多设备并行执行

当 TensorFlow 运行一个 TF 函数时，它会先分析其图，找出需要评估的操作列表，并统计每个操作的依赖数量。然后，将每个没有依赖的操作（即每个源操作）添加到该操作所在设备的评估队列中。操作评估完成后，依赖该操作的每个操作的依赖计数器会减 1。当一个操作的依赖计数器达到 0 时，它会被推送到其所在设备的评估队列中。当所有输出都计算完成后，将返回结果。

以下是 CPU 和 GPU 上操作执行的具体情况：
|设备|执行情况|
|----|----|
|CPU|CPU 评估队列中的操作会被分发到一个名为 inter-op 线程池的线程池中。如果 CPU 有多个核心，这些操作实际上会并行评估。一些操作有多线程 CPU 内核，这些内核会将任务拆分为多个子操作，这些子操作会被放置在另一个评估队列中，并分发到一个名为 intra-op 线程池的第二个线程池中（所有多线程 CPU 内核共享）。简而言之，多个操作和子操作可以在不同的 CPU 核心上并行评估。|
|GPU|GPU 评估队列中的操作会按顺序评估。然而，大多数操作有多线程 GPU 内核，通常由 TensorFlow 依赖的库（如 CUDA 和 cuDNN）实现。这些实现有自己的线程池，通常会尽可能多地利用 GPU 线程（这就是为什么 GPU 不需要 inter-op 线程池的原因：每个操作已经占用了大部分 GPU 线程）。|

graph LR
    classDef startend fill:#F5EBFF,stroke:#BE8FED,stroke-width:2px
    classDef process fill:#E5F6FF,stroke:#73A6FF,stroke-width:2px

    A([开始]):::startend --> B(分析 TF 函数图):::process
    B --> C(找出源操作):::process
    C --> D(将源操作添加到对应设备评估队列):::process
    D --> E{操作评估完成?}:::process
    E -- 是 --> F(减少依赖操作的计数器):::process
    F --> G{计数器为 0?}:::process
    G -- 是 --> H(将操作推送到对应设备评估队列):::process
    H --> I{所有输出计算完成?}:::process
    I -- 是 --> J([返回结果]):::startend
    E -- 否 --> D
    G -- 否 --> E
    I -- 否 --> E

你可以通过调用 tf.config.threading.set_inter_op_parallelism_threads() 控制 inter-op 线程池中的线程数量，通过 tf.config.threading.set_intra_op_parallelism_threads() 设置 intra-op 线程数量。这在你不希望 TensorFlow 使用所有 CPU 核心或希望它单线程运行时很有用。

利用 GPU 的强大功能，你可以进行以下操作：
- 并行训练多个模型，每个模型使用自己的 GPU：为每个模型编写训练脚本并并行运行，设置 CUDA_DEVICE_ORDER 和 CUDA_VISIBLE_DEVICES ，使每个脚本只看到一个 GPU 设备。这对于超参数调优非常有用，因为你可以并行训练多个具有不同超参数的模型。例如，如果你有一台带有两个 GPU 的机器，在一个 GPU 上训练一个模型需要一小时，那么在两个 GPU 上并行训练两个模型也只需要一小时。
- 在单个 GPU 上训练模型，并在 CPU 上并行执行所有预处理操作，使用数据集的 prefetch() 方法提前准备好接下来的几个批次，以便在 GPU 需要时可用。
- 如果你的模型接受两个图像作为输入，并在合并输出之前使用两个 CNN 进行处理，将每个 CNN 放在不同的 GPU 上可能会使运行速度更快。
- 创建一个高效的集成模型：将不同的训练好的模型放在每个 GPU 上，这样可以更快地获得所有预测结果，从而生成集成模型的最终预测。

5. 跨多个设备训练模型

有两种主要方法可以跨多个设备训练单个模型：
- 模型并行 ：将模型拆分为多个部分，每个部分在不同的设备上运行。但这种方法比较复杂，其有效性很大程度上取决于神经网络的架构。对于全连接网络，这种方法通常收益不大。例如，直观上将每个层放在不同的设备上并不可行，因为每个层都需要等待前一层的输出才能进行操作。垂直分割模型（例如，将每层的左半部分放在一个设备上，右半部分放在另一个设备上）稍微好一些，因为每层的两个半部分可以并行工作，但问题是下一层的每个半部分都需要两个半部分的输出，因此会有大量的跨设备通信，这可能会完全抵消并行计算的好处。不过，一些神经网络架构（如卷积神经网络）包含与下层部分连接的层，因此更容易有效地将部分模型分布在不同设备上。深度循环神经网络也可以更有效地跨多个 GPU 拆分，但在实践中，在单个 GPU 上运行常规的 LSTM 层堆栈实际上更快。
- 数据并行 ：将模型复制到每个设备上，每个副本在不同的数据子集上进行训练。这是一种更简单且通常更高效的方法，后续会详细介绍。

使用 GPU 加速 TensorFlow 计算

5. 跨多个设备训练模型（续）

在前面我们提到了跨多个设备训练单个模型有模型并行和数据并行两种主要方法，下面详细介绍数据并行方法。

数据并行

数据并行是将模型复制到每个设备上，每个副本在不同的数据子集上进行训练。这种方法操作相对简单，并且通常效率较高。

具体来说，在数据并行训练中，每个设备上都有一份模型的副本。训练数据会被分割成多个子集，每个子集分配给一个设备上的模型副本进行训练。在每个训练步骤中，每个设备上的模型副本会根据分配给自己的子集数据计算梯度，然后这些梯度会被汇总，用于更新所有设备上的模型副本的参数。

数据并行有同步和异步两种模式：
- 同步数据并行 ：在同步模式下，所有设备上的模型副本会同时进行前向传播和反向传播计算梯度。当所有设备都完成梯度计算后，会汇总这些梯度并更新模型参数。只有在所有设备都完成参数更新后，才会开始下一轮的训练。这种模式可以保证模型参数在每个训练步骤后都是一致的，但可能会受到最慢设备的限制，导致整体训练速度变慢。
- 异步数据并行 ：在异步模式下，每个设备上的模型副本独立地进行训练和参数更新，不需要等待其他设备。每个设备在完成自己的梯度计算后，会立即更新模型参数。这种模式可以充分利用每个设备的计算资源，提高训练速度，但可能会导致模型参数在不同设备上存在一定的差异。

以下是一个简单的同步数据并行训练的示例代码框架：

import tensorflow as tf

# 假设我们有多个 GPU 设备
strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    # 定义模型
    model = tf.keras.Sequential([...])
    # 编译模型
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

# 加载和分割数据
dataset = ...
dist_dataset = strategy.experimental_distribute_dataset(dataset)

# 训练模型
model.fit(dist_dataset, epochs=10)

在这个示例中， tf.distribute.MirroredStrategy() 用于创建一个同步数据并行的策略。 with strategy.scope() 上下文管理器确保模型和优化器在所有设备上进行复制。 strategy.experimental_distribute_dataset() 用于将数据集分割并分发给每个设备。

6. 总结与应用建议

通过前面的介绍，我们了解了如何使用 GPU 加速 TensorFlow 计算，包括检查 GPU 识别、管理 GPU 内存、在设备上放置操作和变量、多设备并行执行以及跨多个设备训练模型等方面。以下是一些应用建议：
- 内存管理 ：
- 如果需要同时运行多个 TensorFlow 程序，根据实际情况选择合适的内存管理方法。例如，如果每个程序需要固定的内存量，可以使用限制 TensorFlow 占用的 GPU 内存的方法；如果内存使用情况不确定，可以考虑按需分配 GPU 内存的方法。
- 在生产环境中，尽量避免使用按需分配内存的方法，因为它可能会导致不确定性的行为。
- 设备放置 ：
- 遵循将数据预处理操作放在 CPU 上，神经网络操作放在 GPU 上的原则，以充分利用 CPU 和 GPU 的优势。
- 避免不必要的数据进出 GPU 传输，减少通信开销。
- 并行执行 ：
- 根据模型的特点和硬件资源，选择合适的并行执行方式。例如，如果模型可以拆分为多个独立的部分，可以考虑将这些部分放在不同的 GPU 上并行执行；如果是训练多个模型，可以为每个模型分配一个独立的 GPU。
- 跨设备训练 ：
- 对于简单的模型或数据量较小的情况，可以先尝试数据并行方法，因为它相对简单且高效。
- 如果模型架构适合，并且有足够的硬件资源和技术能力，可以考虑模型并行方法，但需要注意跨设备通信的问题。

graph LR
    classDef startend fill:#F5EBFF,stroke:#BE8FED,stroke-width:2px
    classDef process fill:#E5F6FF,stroke:#73A6FF,stroke-width:2px

    A([开始训练]):::startend --> B{选择训练方法}:::process
    B -- 模型并行 --> C(拆分模型到不同设备):::process
    B -- 数据并行 --> D(复制模型到每个设备):::process
    C --> E(各设备计算梯度):::process
    D --> F(分割数据到各设备):::process
    F --> E
    E --> G(汇总梯度):::process
    G --> H(更新模型参数):::process
    H --> I{训练完成?}:::process
    I -- 否 --> E
    I -- 是 --> J([结束训练]):::startend

总之，合理利用 GPU 资源可以显著提高 TensorFlow 模型的训练和推理速度。通过掌握上述方法和技巧，并根据实际情况进行调整和优化，能够更好地发挥 GPU 的性能，加速深度学习任务的完成。

以下是一个简单的表格总结不同训练方法的特点：
|训练方法|优点|缺点|适用场景|
|----|----|----|----|
|模型并行|可利用多设备资源处理大型模型|实现复杂，跨设备通信开销大|适合特定架构的大型模型|
|数据并行|操作简单，效率较高|可能受最慢设备限制（同步模式）|大多数深度学习模型训练场景|