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

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 "simulated 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 three-dimensional (3D) low-bit quantization on modern graphics processing units (GPUs) is still unexplored.

We introduce MedPTQ, an open-source pipeline for real post-training quantization that implements true 8-bit (INT8) inference on state-of-the-art (SOTA) 3D medical segmentation models, i.e., U-Net, SegResNet, SwinUNETR, nnU-Net, UNesT, TransUNet, ST-UNet, and VISTA3D. MedPTQ involves two main steps. First, we use TensorRT to perform simulated quantization for both weights and activations with an unlabeled calibration dataset. Second, we convert this simulated quantization into real quantization via the TensorRT engine on real GPUs, resulting in real-world reductions in model size and inference latency.

Extensive experiments benchmark MedPTQ across seven models and three datasets and demonstrate that it effectively performs INT8 quantization on GPUs, reducing model size by up to 3.83× and latency by up to 2.74×, while maintaining nearly identical Dice similarity coefficient (mDSC) performance to FP32 models. This advancement enables the deployment of efficient deep learning models in medical imaging applications where computational resources are constrained. The MedPTQ 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.

We have introduced MedPTQ, a real post-training quantization pipeline that delivers INT8 inference for SOTA 3D artificial intelligence (AI) models in medical imaging segmentation. MedPTQ effectively reduces real-world model size, computational requirements, and inference latency without compromising segmentation accuracy on modern GPUs, as evidenced by mDSC comparable to full-precision baselines. We validate MedPTQ across a diverse set of AI architectures, ranging from convolutional-neural-network-based to transformer-based models, and a wide variety of medical imaging datasets. These datasets are collected from multiple hospitals with distinct imaging protocols, cover different body regions (such as the brain, abdomen, or full body), and include multiple imaging modalities [computed tomography (CT) and magnetic resonance imaging (MRI)]. Collectively, these results highlight our MedPTQ's strong generalizability and adaptability for a broad spectrum of medical imaging tasks.