如何选择/配置 AWS GPU 实例以加速 TensorFlow.keras？

Question

我有一个包含大约 600MB 训练数据的 LSTM tf.keras 模型。每个训练 epoch 大约需要 90 秒。我有最新版本的 tensorflow，即 v2.2。它在 AWS g3.4xlarge 实例上运行。此实例具有 Nvidia 的 Tesla M60 GPU，并为 GPU 提供 8GB RAM。

我想做超参数调优，所以我需要加快执行速度。因此，我将模型和数据移动到了一个 AWS p3.2xlarge 实例，该实例有一个 P100 GPU 和 16GB 的 RAM。然后我发现每个 epoch 的训练时间根本没有变化。

所以我切换到更大的 AWS 实例 p3.8xlarge，它有 4 个 Tesla V100 GPU 和总共 64GB 的 RAM。在第一次运行中，它只使用了 1 个 GPU，并且每个 epoch 的执行时间相同，大约 90 秒。于是我在tensorflow网站上找到一篇关于如何使用所有GPU的文章，https://www.tensorflow.org/guide/gpu

因此，在所有 4 个 GPU 都运行的情况下，一个 epoch 的执行时间从 90 秒变为 112 秒！我使用 nvtop 和 nvidia-smi 来监控 GPU 性能，如下图所示。

我需要做些什么来减少执行时间？

Answer 1

在您不更改网络中的任何内容之前，运行每个 epoch 的时间显然是相同的。仅仅在更大的 GPU 上运行它并不是解决问题的方法。

首先，如果您可以更改网络以减少参数数量，那就太好了。所以，减少你的模型显然是让它运行得更快的第一件事。

但是，如果这是不可能的，你可以做两件事：

1. 使用 tf.mixed_precision 更快地运行您的模型。

from tensorflow.keras.mixed_precision import experimental as mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

它以半精度格式执行运算，同时以单精度存储最少的信息，从而在网络的关键部分保留尽可能多的信息，从而显着提高计算速度。

1. 使用 XLA。

import tensorflow as tf

tf.config.optimizer.set_jit(True)

加速线性代数，XLA 是用于矩阵运算的特定领域编译器。它可以使您的网络更快，而无需更改您的源代码。

请尝试这两种方法。我个人使用过混合精度，它确实减少了网络时间。

另外，请下次不要跳到更大的实例，因为这是浪费你的钱。只需尝试减少参数数量（即网络大小）或使用这两个技巧。如果我发现任何新技巧，我会更新这个答案。