MedPTQ: a practical pipeline for real post-training quantization in 3D medical image segmentation

Area of Science:

• Medical Imaging AI
• Deep Learning Optimization
• Computational Efficiency

Background:

• Quantizing deep neural networks reduces computational demands, crucial for medical imaging.
• Existing methods often simulate quantization, failing to reduce model size or improve real-world speed.
• Real 3D low-bit quantization on modern GPUs remains largely unexplored.

Purpose of the Study:

• Introduce MedPTQ, an open-source pipeline for real post-training quantization.
• Implement true 8-bit (INT8) inference for state-of-the-art 3D medical segmentation models.
• Enable efficient deployment of AI in resource-limited medical imaging settings.

Main Methods:

• Utilize TensorRT for simulated quantization of weights and activations with an unlabeled dataset.
• Convert simulated quantization to real INT8 quantization using TensorRT engine on GPUs.
• Apply to various 3D medical segmentation models including U-Net, SwinUNETR, and nnU-Net.

Main Results:

• MedPTQ achieves real INT8 quantization on GPUs, reducing model size up to 3.83× and latency up to 2.74×.
• Maintains nearly identical Dice similarity coefficient (mDSC) performance compared to full-precision (FP32) models.
• Demonstrated effectiveness across seven models and three diverse datasets.

Conclusions:

• MedPTQ provides real INT8 inference for 3D AI medical imaging segmentation models.
• Reduces model size, computational needs, and latency on GPUs without compromising accuracy.
• Shows strong generalizability across diverse AI architectures, datasets, and imaging modalities.