AI模型推理的最低系统要求

Question

如何确定运行 pytorch 模型推理以避免内存不足 (OOM) 的最低系统要求？

例如，我有一个经过微调的文本分类模型（“sentiment_model.pt”）和一个从该 .pt 文件加载权重并使用它进行推理的 Python 代码。当我在 8GB 机器上运行代码时，它会因为 OOM 而被终止，而在 16GB 机器上它运行得很好。我可以优化代码以在 8GB 机器上运行吗？如何确定最低系统要求？

Answer 1

这很难回答，因为这在很大程度上取决于模型和您使用的批量大小。为了进行推理，pytorch 将在内存中设置所有参数和所有潜在张量。你需要评估这两个东西的大小才能知道内存使用情况。

例如，CNN 8 层（3 x 3 卷积）在输入中采用 256 x 256 图像，其中每层将潜在图的大小减少 2，并将深度乘以 2。如果

d

是您的第一个深度大致有

9*(2⁷+2⁸)*d²=3456*d²

个参数，最大的潜在变量的大小为

2^16*d*batch_size=16384*d*batch_size

。

那么如果您使用 float32（4 字节/数字），大小将约为

65kB*d*batch_size + 13kB*d²

。然后，对于这种 CNN，您可以看到，您需要严格的最小

13kB*d²

内存，然后当您将批量大小增加到 d/4 以上时，需要显着更多的内存。

有时，即使批量为 1，主导项也是批量大小的项，有时它始终是所有合理批量大小的常数项。