arXiv:2501.17343cs.CVcs.AI2025-01被引 12

实现3D医学图像分割模型的真8位量化,显著压缩模型且不损失性能。

Post-Training Quantization for 3D Medical Image Segmentation: A Practical Study on Real Inference Engines

  • 用TensorRT先假量化再转为真实8位量化,适配现代GPU。
  • 在6个主流3D分割模型上实现无损8位量化,推理速度提升明显。
  • 适合资源受限场景下的医疗影像部署,代码开源可复现。

降低深度神经网络计算精度(位宽)可显著减少内存占用并加速推理,使模型更适用于计算资源受限的大规模医学影像应用。然而,现有方法多基于‘假量化’,仅模拟低精度运算,无法实际缩减模型尺寸或提升真实推理速度。此外,现代GPU上实现真正3D低比特量化的潜力尚未探索。本研究提出一种真实后训练量化(PTQ)框架,成功在SOTA 3D医学分割模型(U-Net、SegResNet、SwinUNETR、nnU-Net、UNesT、TransUNet、ST-UNet和VISTA3D)上实现真正的8位量化。方法分两步:首先使用TensorRT对权重与激活进行假量化,基于无标签校准数据集;其次通过TensorRT引擎在真实GPU上转换为真量化,实现模型尺寸与推理延迟的真实下降。大量实验表明,该框架在不牺牲模型性能的前提下,有效实现GPU上的8位量化。此进展推动了高效深度学习模型在资源受限医疗影像场景中的部署。代码与预训练模型已发布,涵盖BTCV(腹部13标签)、Whole Brain Dataset(全脑133标签)及TotalSegmentator V2(全身体104标签)上的多个模型。项目地址:https://github.com/hrlblab/PTQ。

原文摘要 · Abstract (English)

Quantizing deep neural networks ,reducing the precision (bit-width) of their computations, can remarkably decrease memory usage and accelerate processing, making these models more suitable for large-scale medical imaging applications with limited computational resources. However, many existing methods studied "fake quantization", which simulates lower precision operations during inference, but does not actually reduce model size or improve real-world inference speed. Moreover, the potential of deploying real 3D low-bit quantization on modern GPUs is still unexplored. In this study, we introduce a real post-training quantization (PTQ) framework that successfully implements true 8-bit quantization on state-of-the-art (SOTA) 3D medical segmentation models, i.e., U-Net, SegResNet, SwinUNETR, nnU-Net, UNesT, TransUNet, ST-UNet,and VISTA3D. Our approach involves two main steps. First, we use TensorRT to perform fake quantization for both weights and activations with unlabeled calibration dataset. Second, we convert this fake quantization into real quantization via TensorRT engine on real GPUs, resulting in real-world reductions in model size and inference latency. Extensive experiments demonstrate that our framework effectively performs 8-bit quantization on GPUs without sacrificing model performance. This advancement enables the deployment of efficient deep learning models in medical imaging applications where computational resources are constrained. The code and models have been released, including U-Net, TransUNet pretrained on the BTCV dataset for abdominal (13-label) segmentation, UNesT pretrained on the Whole Brain Dataset for whole brain (133-label) segmentation, and nnU-Net, SegResNet, SwinUNETR and VISTA3D pretrained on TotalSegmentator V2 for full body (104-label) segmentation. https://github.com/hrlblab/PTQ.

3D分割量化医学影像TensorRT

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。