是否可以为混合精度 NN 设置 GPU 亲和力，FP32 和 FP16 去往不同的 GPU？

Posted 2023-04-15

【中文标题】是否可以为混合精度 NN 设置 GPU 亲和力，FP32 和 FP16 去往不同的 GPU？

【英文标题】：Is it possible to set GPU affinity for a mixed precision NN, with FP32 and FP16 going to different GPUs?

【发布时间】：2019-07-10 13:13:27

【问题描述】：

我有一个 GTX 1080 和一个 RTX 2080。我想同时使用两者进行训练，但由于 RTX 处理 FP16 的速度是原来的两倍，我想将其设置为多 GPU 和 RTX处理 FP16 层，GTX 处理 FP32 层。

这在 tensorflow、pytorch 或 keras 下是否可行？

【参考方案1】：

张量流

在 TF 中，可以为每个层指定要在哪个设备上执行（GPU、CPU 或特定 GPU，如果您有多个 GPU ...）。这是使用with tf.device('device_name') 语句完成的（您需要提供有意义的device_name）。请参阅Using multiple GPUs 部分。

Keras

既然这在 TF 中是可能的，这意味着你也可以在 Keras 中使用它，如果你使用 TF 作为 Keras 后端（Keras 只是一个高级神经网络 API）。

请注意，在 Keras 中有一个 multi_gpu_model() 函数在 Keras 中，但它只在多个 GPU 上复制整个模型，您无法指定在特定 GPU 上放置哪个层。