[論文レビュー] 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 本ツールボックスは、光シート顕微鏡法を用いたオステオサイトのセマンティックセグメンテーションといった大規模3次元バイオイメージ解析タスクをどの程度高速化できるか？
• RQ4 微細な構造的詳細が重要なラベルフリー蛍光画像予測タスクにおいて、圧縮モデルが忠実な再現性を維持できるか？
• RQ5 本ツールボックスは、mmv_im2im や U-Net といった広く使われているバイオイメージングAIフレームワークとどの程度互換性があり、実際の応用でどの程度のパフォーマンス向上が達成できるか？

主な発見

• 4つの異なるバイオイメージ解析タスク（3次元セグメンテーションやラベルフリーイメージングを含む）において、推論時間が2–5倍速くなった。
• 推論中、エネルギー消費が30–80%削減され、AIベースの解析の環境的影響が顕著に低減された。
• 一部のケースでは、過学習の低減により圧縮後に精度が向上した。これは、スプライシングと量子化が正則化効果を示していることを示している。
• 視覚的検査および人間によるアノテーションの正解との定量的比較により、圧縮モデルは元のモデルとほぼ同一の予測品質を維持していた。
• 2分の推論タスクが近似リアルタイムにまで高速化され、迅速な反復的メソッド開発や大規模スクリーニングが可能になった。
• 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.).