如何仅花 25 美元并在 3 小时内完成 ImageNet 训练？

2023-12-13 16:07:00

DAWNBench 是斯坦福大学建立的项目，旨在以竞赛形式对不同深度学习方法加以比较。该项目要求参赛者使用固定的硬件和软件环境，在有限的计算资源和时间内完成特定的深度学习任务。

在 DAWNBench 竞赛当中，CIFAR10 与 ImageNet 的表现引起了我们的关注。二者的目标基本一致，在保证相对合理准确度（ImageNet 为 93%，CIFAR 为 94%）的前提下，尽量减少训练时间和计算成本。CIFAR10 是一个包含 60,000 张图像的数据集，分为 10 个类别。ImageNet 是一个包含超过 1000 万张图像的数据集，分为 1000 个类别。

CIFAR10 的训练相对简单，可以在普通的个人电脑上完成。ImageNet 的训练则需要大量的数据和计算资源。在 DAWNBench 竞赛中，ImageNet 的训练被限制在 3 个小时以内，计算资源也被限制在单个 NVIDIA Tesla V100 显卡。

为了在有限的计算资源和时间内完成 ImageNet 的训练，我们采用了以下策略：

使用预训练模型。预训练模型是指已经在其他数据集上训练过的模型。我们使用了一个在 ImageNet 上预训练的模型作为 ImageNet 训练的起点。这可以帮助我们快速收敛，并提高训练的准确性。
使用数据增强技术。数据增强技术是指通过对训练数据进行随机变换（如裁剪、翻转、旋转等）来增加训练数据的数量和多样性。这可以帮助模型学习到更 general 的特征，并提高模型的泛化能力。
使用梯度累积。梯度累积是一种训练技巧，可以减少对显存的要求。它将多个 batch 的梯度累积起来，然后才更新模型参数。这可以让我们使用更大的 batch size，从而提高训练速度。
使用混合精度训练。混合精度训练是一种训练技巧，可以减少对显存的要求。它将模型的权重和激活值存储在不同的精度上。这可以让我们使用更大的模型，从而提高训练的准确性。
使用分布式训练。分布式训练是一种训练技巧，可以利用多个 GPU 并行训练模型。这可以大幅提高训练速度。
使用高性能计算资源。高性能计算资源是指具有强大计算能力的计算设备。我们使用了一个 NVIDIA DGX-1 系统来完成 ImageNet 的训练。这个系统包含 8 个 NVIDIA Tesla V100 显卡，总算力达到 170 TFLOPS。