边缘设备上部署模型的限制之一——显存占用：模型的参数量只是冰山一角

边缘设备上部署模型的限制之一——显存占用：模型的参数量只是冰山一角

在边缘设备上部署深度学习模型已成为趋势，但资源限制是其核心挑战之一。其中，显存（或更广义的内存）占用是开发者们必须仔细考量的重要因素。许多人认为显存占用主要取决于模型的参数量，这种看法虽然没错，但并不全面。实际上，显存的占用远不止模型参数量那么简单。

关于边缘设备（Edge Device）的介绍，可以参见我的这一篇文章：Edge Device（边缘设备）：连接物理世界与数字世界的桥梁

模型参数量：显存占用的基石

首先，我们不得不承认，模型参数量确实是显存占用的一个主要且直接的决定因素。

• 模型权重存储：神经网络中的每一个权重和偏置都需要存储在显存中。模型的参数越多，所需存储的空间就越大。
• 数据类型影响：参数的数据类型也至关重要。例如，使用单精度浮点数（FP32）存储每个参数需要4字节；而如果采用半精度（FP16）或8位整型量化（INT8），则分别只需要2字节或1字节。通过数据类型优化，即使参数量不变，也能显著降低显存需求。

---

显存占用：远不止参数量那么简单

然而，如果我们只关注参数量，就会错过显存占用的其他几个关键组成部分，这些部分有时甚至比参数量本身占用更大的显存。

1. 中间激活（Intermediate Activations）

在模型推理过程中，数据流经每一层网络时会产生中间结果，我们称之为激活（Activations）或特征图（Feature Maps）。这些中间激活需要临时存储在显存中，以便后续层进行计算。

• 批处理大小（Batch Size）：批处理大小直接影响中间激活的显存占用。批次越大，每次推理需要处理的样本越多，每层生成的特征图尺寸就越大，从而占据更多的显存。
• 网络深度与宽度：模型的层数（深度）和每层的神经元数量或特征图通道数（宽度）都会直接影响中间激活的存储需求。更深或更宽的网络往往会产生海量的中间激活。

2. 优化器状态（Optimizer State）

这一点主要与模型的训练而非推理有关。在使用Adam、RMSprop等优化器进行模型训练时，这些优化器会为每个模型参数维护额外的状态变量（例如，动量和方差的估计值）。这些状态变量通常与参数数量相同，并且同样需要存储在显存中。在某些情况下，优化器状态的显存占用甚至可能使总显存需求翻倍或更多。

3. 输入和输出数据

模型在进行推理之前，输入数据（如图像、音频、传感器数据等）需要加载到显存中。输入数据越大（例如高分辨率图像），占用的显存就越多。同样，模型的最终输出结果也需要一定的显存来存储，尽管这部分占用通常相对较小。

4. 辅助数据结构与操作开销

除了上述几点，还有一些其他因素也会贡献显存占用：

• 模型图结构：深度学习框架（如TensorFlow Lite、PyTorch Mobile）在加载模型时，需要存储模型的计算图结构信息，这也会占用少量显存。
• 缓冲区/缓存：某些特定的计算操作可能需要额外的缓冲区来进行数据转换或中间计算。
• 量化和剪枝的潜在开销：尽管量化和剪枝旨在减少模型大小，但在某些实现中，它们可能在处理过程中引入临时的额外存储开销。

---

总结与思考

综上所述，显存占用是一个多维度的问题。一个拥有高参数量的模型，如果经过了高效的量化（Quantization）和剪枝（Pruning）等优化技术，并且在推理时采用较小的批处理大小，其最终的显存占用可能反而低于一个参数量较少但未经优化、且使用大批处理大小的模型。

在边缘设备上部署深度学习模型时，我们必须跳出“只看参数量”的误区，全面评估和优化显存占用。这包括对模型架构的选择、数据类型的量化、批处理大小的设定以及中间激活的管理等方面进行深入考量，以确保模型能在有限的资源下稳定、高效地运行。

希望这篇文章能帮助你更好地理解边缘设备上显存占用的构成！