[论文解读] EfficientBioAI: Making Bioimaging AI Models Efficient in Energy, Latency and Representation

Yu Zhou, Justin Sonneck+5 more|arXiv (Cornell University)|Jun 9, 2023

Cell Image Analysis TechniquesBiochemistry, Genetics and Molecular Biology被引用 3

一句话总结

EfficientBioAI 是一个即插即用的工具箱，通过剪枝和量化技术压缩生物影像AI模型，分别将能耗降低30–80%、推理延迟减少2–5倍，且不损失精度——有时甚至通过减少过拟合现象而提升精度。该工具与PyTorch、U-Net和mmv_im2im等常见生物影像分析框架兼容，可在各类生物医学影像任务中实现CPU和GPU上的高效部署。

ABSTRACT

Artificial intelligence (AI) has been widely used in bioimage image analysis nowadays, but the efficiency of AI models, like the energy consumption and latency is not ignorable due to the growing model size and complexity, as well as the fast-growing analysis needs in modern biomedical studies. Like we can compress large images for efficient storage and sharing, we can also compress the AI models for efficient applications and deployment. In this work, we present EfficientBioAI, a plug-and-play toolbox that can compress given bioimaging AI models for them to run with significantly reduced energy cost and inference time on both CPU and GPU, without compromise on accuracy. In some cases, the prediction accuracy could even increase after compression, since the compression procedure could remove redundant information in the model representation and therefore reduce over-fitting. From four different bioimage analysis applications, we observed around 2-5 times speed-up during inference and 30-80$\%$ saving in energy. Cutting the runtime of large scale bioimage analysis from days to hours or getting a two-minutes bioimaging AI model inference done in near real-time will open new doors for method development and biomedical discoveries. We hope our toolbox will facilitate resource-constrained bioimaging AI and accelerate large-scale AI-based quantitative biological studies in an eco-friendly way, as well as stimulate further research on the efficiency of bioimaging AI.

研究动机与目标

• 应对日益复杂的生物影像AI模型在资源受限及大规模生物学研究中带来的能耗与延迟成本上升问题。
• 解决缺乏针对生物学家及现有生物影像AI框架量身定制的、易于使用的通用模型压缩工具的问题。
• 在不牺牲预测性能的前提下，实现在CPU、边缘设备及多种硬件后端上的高效AI模型部署。
• 证明模型压缩可通过移除冗余表征减少过拟合，从而潜在提升模型泛化能力。
• 通过广泛提供高效AI模型，推动环保、可扩展的AI驱动生物发现。

提出的方法

• 将神经网络剪枝与量化作为核心压缩技术，以减小模型大小并降低计算负载。
• 应用训练后量化，将32位浮点权重与激活值转换为8位整数，最大限度减少内存访问与能耗。
• 采用结构化剪枝移除低重要性的卷积核或神经元，降低模型复杂度，同时保留特征表示能力。
• 支持与基于PyTorch的模型及CellPose、DeepImageJ和mmv_im2im等常见生物影像工具的即插即用集成。
• 将多种压缩策略（如先剪枝后量化）组合为可自定义的流水线，实现最优效率。
• 通过优化的推理引擎在CPU和GPU平台上部署压缩模型，确保跨平台兼容性与低延迟推理。

Figure 1: overview of the toolbox. EfficientBioAI aims to compress the model and accelerate the bioimage tasks. a. Given a pretrained neural network (NN) in pytorch, the whole pipeline consists of two phases: compression (pruning, quantization) and inference. The compressed models are able to run on

实验结果

研究问题

• RQ1 通过剪枝与量化实现的模型压缩，是否能在不降低精度的前提下显著减少生物影像AI模型的能耗与推理延迟？
• RQ2 在数据量有限的情况下，量化与剪枝等压缩技术在多大程度上可通过减少过拟合来提升模型泛化能力？
• RQ3 该工具箱在加速大规模3D生物影像分析任务（如光片显微镜中骨细胞的语义分割）方面的实际效果如何？
• RQ4 在标签无荧光图像预测任务中，压缩后的模型是否能保持对细微结构特征的高保真度？
• RQ5 该工具箱与mmv_im2im和U-Net等广泛使用的现有生物影像AI框架的兼容性如何？实际应用中可实现多大的性能提升？

主要发现

• 在四项不同的生物影像分析任务中，包括3D分割与无标签成像，该工具箱实现了2–5倍的推理速度提升。
• 推理过程中的能耗降低了30–80%，显著降低了AI分析的环境影响。
• 在某些情况下，压缩后模型精度反而提升，归因于剪枝与量化带来的正则化效应，减少了过拟合。
• 压缩模型与原始模型在预测质量上几乎完全一致，经人工标注的真实值对比验证，视觉检查与定量分析均表明其性能保持稳定。
• 该工具箱成功将原本需两分钟的推理任务加速至接近实时，支持更快的迭代方法开发与大规模筛选。
• 该方法与mmv_im2im、CellPose和U-Net等流行生物影像AI工具表现出强兼容性，展现出在该领域广泛适用的潜力。

Figure 2: The effect of the sparsity ratio on different pruners, i.e., pruning algorithms (L1, L2, and FPGM) by observing the ratio of performance degradation (a.), energy saving (b.) and latency speedup (c.).