Lightweight nnU-Net with Knowledge Distillation and Multi-Threaded Optimization for Abdominal Organ Segmentation

Houkai, fanggang, Yue Liu, luoyinyin, huangbingquan, tangbin

Keywords: Knowledge Distillation, epthwise Separable Convolution, Pyramid Pooling, Multi-threading Optimization, Bottleneck Residual Blocks

TL;DR: Model improvement based on nnunet model on 3dCT dataset

Abstract: Although current deep learning models have achieved remarkable success in medical image segmentation, their deployment on resource-constrained environments remains challenging due to substantial computational and memory requirements, particularly for 3D medical images. Existing lightweight models often sacrifice segmentation accuracy significantly to reduce computational overhead.To address this challenge, we propose a comprehensive lightweight optimization framework based on nnU-Net that maintains high segmentation accuracy while dramatically reducing computational requirements. Our main contributions include: (1) a lightweight network architecture that replaces standard 3D convolutions with depthwise separable convolutions and incorporates bottleneck residual blocks, reducing model parameters to 157.59K while preserving the feature representation capability; (2) pyramid pooling modules for enhanced multi-scale feature extraction and improved boundary precision; (3) a knowledge distillation strategy where a teacher network (original nnU-Net) transfers knowledge to our lightweight student network through feature-level and output-level distillation losses; and (4) a multi-threaded inference optimization system that parallelizes post-processing operations using 12 threads, achieving 2-4× speedup in post-processing.Comprehensive experiments on MICCAI FLARE 2025 validation set validate the effectiveness of our approach. The proposed method achieves an average organ Dice Similarity Coefficient (DSC) of 90.48\% and Normalized Surface Dice (NSD) of 96.56\%, and the validation score improved by 0.9 points. Our method ranked 1st in the MICCAI FLARE 2025 Task 2 online validation leaderboard. The code is available at:https://github.com/houkainiubi/flare25task2_hk.git

Submission Number: 12