批次太大：找到适合GPU的批次大小

一个简单的功能，可以识别可以填充GPU内存的Pytorch模型的批处理大小

我敢肯定，你们中的许多人都有以下痛苦的经历：您在GPU上进行了多次ML实验，以训练一夜之间，当您回来10小时后检查时，您意识到，由于硬件未充分利用或更糟糕，进度栏几乎没有移动，所有实验由于内存（OOM）误差而失败。在这个迷你指南中，我们将实施一种自动化方法，以找到可以充分利用GPU内存的Pytorch模型的批处理大小而不会引起OOM！

除了模型体系结构和参数数量之外，批处理大小是控制实验使用的GPU内存量的最有效的超参数。找到可以充分利用加速器的最佳批次大小的适当方法是通过GPU分析，这是一种在计算设备上监视进程的过程。两个都TensorFlow和Pytorch提供有关如何在其框架中进行分析的详细指南和教程。此外，批处理大小会极大地影响模型的性能。例如，大批量大小会导致概括不佳，请查看此博客文章凯文·沉在批次大小对训练动态的影响如果您对此主题感兴趣。

但是，如果您只想训练模型来测试想法，分析或执行超参数搜索以找到最佳批次大小可能是过高的杀伤力，尤其是在项目的早期阶段。找到可让您在没有OOM的情况下适合模型的价值的常见方法是使用诸如批量尺寸的小批量训练模型，同时使用诸如诸如nvidia-smi要么nvitop。然后，如果模型不足，则增加值，并重复该过程，直到击中内存容量为止。但是，此手动过程可能很耗时。更烦人的是，当您必须在不同的GPU上运行具有不同内存大小的不同GPU实验时，则必须重复每个设备的相同过程。幸运的是，我们可以将这个乏味的迭代过程转换为代码，并在实际实验之前运行它，以便您知道自己的模型不会引起OOM。

这个想法非常简单：

初始化模型。
将批量大小设置为2（用于批处理）
创建具有样本形状作为真实数据的虚拟数据。
训练模型n步骤（向前和后退）。
如果模型没有错误运行，请增加批处理大小，然后转到步骤3。如果OOM升高（即Runtimeerror在pytorch中）然后将批量大小设置为上一个值并终止。
返回最终批次大小。

将其放入代码中

如您所见，此功能有7个参数：

模型- 您要适合的模型，请注意，该模型将在功能末尾从内存中删除。
设备-火炬应该是CUDA设备。
input_shape- 数据的输入形状。
output_shape- 模型的预期输出形状。
dataset_size- 数据集的大小（当批处理大小已经大于数据集的大小时，我们不想继续进行搜索）。
max_batch_size- 可选的参数，以设置要使用的最大批量大小。
num_iterations- 在增加批处理大小之前更新模型的迭代次数，默认为5。

让我们快速浏览功能中发生的事情。我们首先将模型加载到GPU，初始化Adam Optimizer，然后将初始批次大小设置为2（如果不使用batchNorm，则可以从批处理大小为1开始）。然后，我们可以开始迭代过程。首先，我们检查当前批处理大小是否大于数据集的大小或最大所需批处理大小，如果是的话，我们会破坏循环。否则，我们创建虚拟输入和目标，将其移至GPU并适合模型。我们训练该模型5个步骤，以确保向前也没有向后传球会导致OOM。如果一切都很好，我们将批量大小乘以2并重新构建模型。如果在上述步骤中发生OOM，则我们将批次大小减少2倍，然后退出环路。最后，我们从内存中清除模型和优化器，然后返回最终批次大小。而已！

请注意，如果OOM的情况，可以继续搜索最佳值（即二进制搜索批处理大小，而是将批处理大小设置为中间点，而不是简单地将批量大小除以2，并继续步骤3。）找到非常适合GPU的批量尺寸。但是，请记住，Pytorch/Tensorflow或其他过程可能会在实验中要求更多的GPU内存，并且您冒险OM，因此我更喜欢有一些摆动的空间。

现在，让我们使用此功能。在这里我们适合RESNET50在1,000列火车的合成图像（3，224，224）生成Fakedata数据集。简而言之，我们首先打电话get_batch_size =（model = resnet（），input_shape = image_shape，output_shape =（num_classes，），dataset_size = dataset_size）要获得可以充分填充GPU内存的批次大小。然后，我们可以初始化模型和数据加载程序，并像平常一样训练模型！

下面的GIF是在NVIDIA RTX 2080 8GB上运行示例代码的屏幕记录。我在find_batch_size功能以显示其正在测试的批次大小，并注意随着功能增加批次大小而增加GPU内存使用情况的增加。我们的脚本确定，批量的大小为16，将导致OOM并以批次大小为8的其余部分，而GPU存储器利用率约为66.8％。

在NVIDIA RTX 2080 8GB上运行示例代码，该脚本确定批次大小为16会导致OOM。[作者的gif]

当我们在NVIDIA RTX 2080TI 11GB上运行完全相同的代码时，我们能够以16个批次大小和GPU内存利用率为90.3％。