Medical Image Segmentation Method Based on Heterogeneous Distillation Strategy
By employing a heterogeneous distillation strategy, dynamic dimension mapping, and multi-level loss weighted fusion, the heterogeneity problem of teacher and student model architectures in medical image segmentation is solved. This enables efficient transfer of anatomical structure knowledge, improves segmentation accuracy and deployment efficiency, and is suitable for resource-constrained clinical equipment.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- CHENGDU UNIV OF TRADITIONAL CHINESE MEDICINE
- Filing Date
- 2026-01-28
- Publication Date
- 2026-04-21
AI Technical Summary
Existing technologies in medical image segmentation suffer from feature space incompatibility issues caused by the heterogeneity of teacher and student model architectures. This makes it difficult to effectively transfer anatomical knowledge, resulting in low segmentation accuracy and low deployment efficiency, which fails to meet clinical needs.
A heterogeneous distillation strategy is adopted, and an adaptive projection matrix is constructed through a dynamic dimension mapping algorithm. Combined with pixel memory and region memory, a five-dimensional loss weighted fusion system is designed to realize multi-level knowledge transfer from the teacher model to the student model. This system includes foreground-sensitive segmentation loss, output layer knowledge distillation loss, intermediate layer knowledge distillation loss, mini-batch pixel contrast loss, and memory contrast loss.
It realizes the transfer of anatomical prior knowledge from complex teacher models to lightweight student models, adapts to resource-constrained clinical deployment scenarios, improves segmentation accuracy while maintaining high efficiency, and increases the inference speed of the student model by 2.3 times on high-configuration devices, meeting the real-time requirements of mobile devices, and maintaining high segmentation accuracy in complex clinical scenarios.
Smart Images

Figure CN121600005B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the interdisciplinary field of medical image processing, model compression, and computer vision. More specifically, this invention relates to a medical image segmentation method based on a heterogeneous distillation strategy. Background Technology
[0002] Medical image segmentation models need to balance accuracy and deployment efficiency. High-performance models have a large number of parameters and high computational costs, making them difficult to deploy on resource-constrained clinical equipment. Knowledge distillation, as a model compression technique, can transfer the "knowledge" of a large model to a smaller model, but traditional distillation methods have significant drawbacks in medical image segmentation.
[0003] Traditional knowledge distillation methods were originally designed for classification tasks, focusing only on feature matching in the output or intermediate layers, neglecting key spatial structural information in medical images (such as the relative position and boundary relationships between lesions and normal tissues). As a result, student models struggle to learn the anatomical structure modeling capabilities of teacher models, resulting in low segmentation accuracy in areas with blurred boundaries and small lesions, which fails to meet clinical requirements.
[0004] While recent structurally-aware distillation methods have introduced attention masks and edge-aware loss to improve the structural awareness capabilities of student models, they have not been deeply integrated with the medical adaptation enhancement module of the teacher model, limiting the quality of knowledge transferred through distillation. Furthermore, the heterogeneous architectures and incompatible feature spaces between the teacher and student models further reduce knowledge transfer efficiency, making it impossible to balance segmentation accuracy and deployment efficiency. Specifically, existing technologies mainly suffer from the following problems:
[0005] 1. Traditional projection matrices are all based on "preset subspaces" (such as linear subspaces spanned by fixed basis vectors), which are only suitable for static alignment of features with the same architecture and the same dimension. When the teacher is a multi-branch model and the student is a single-path model, the feature space misalignment leads to a transfer efficiency of almost 0. Although recent structure-aware distillation focuses on structural features, it has not designed a heterogeneous adaptation module and directly matches the features of heterogeneous models. The problem of incompatibility of intermediate layer feature dimensions has not been solved.
[0006] 2. Traditional distillation relies solely on the comparison of teacher and student features within a single batch during training, resulting in a "time-limited" knowledge transfer. Student models cannot learn the long-term modeling logic of anatomical structures from teachers. While recent structure-aware distillation has introduced structure-related constraints such as attention masks and edge losses, it has not established a knowledge storage mechanism and can only capture local structural features of the current batch, resulting in insufficient knowledge transfer for key areas such as small lesions and fuzzy boundaries.
[0007] 3. Traditional distillation focuses only on the output layer (such as KL divergence matching probability distribution) or a single-level intermediate layer, ignoring key knowledge such as "pixel-level fine-grained features" and "region-level spatial correlation" required by medical images. This results in student models being able to learn only high-level semantics and failing to capture structural details. Although recent structure-aware distillation has added structure-related constraints such as edge loss and attention loss, it has not formed a multi-dimensional fusion system. Knowledge transfer is still concentrated on the single dimension of "structural features" and lacks coverage of difficult-to-classify pixels and historical knowledge.
[0008] 4. While traditional distillation can improve deployment efficiency through model compression, it often comes at the cost of sacrificing segmentation accuracy (such as increased false negative rate for small lesions), failing to meet the core clinical need for "high accuracy". Recently, structure-aware distillation has not optimized the architecture deployment friendliness of student models. Even with lightweight models, the complex feature processing flow still leads to low memory access efficiency, making it difficult to adapt to resource-constrained devices. Summary of the Invention
[0009] One object of the present invention is to solve at least the above-mentioned problems and / or defects, and to provide at least the advantages described below.
[0010] To achieve these objectives and other advantages of the present invention, a medical image segmentation method based on a heterogeneous distillation strategy is provided, comprising:
[0011] S1. Construct a distillation strategy that complements the teacher model and the student model. The distillation strategy includes: a heterogeneous architecture adaptation module based on a dynamic dimension mapping algorithm, a memory enhancement module based on pixel memory bank and region memory bank, and a multi-level knowledge transfer module based on a five-dimensional loss weighted fusion system.
[0012] S2. Input the preprocessed training set into the teacher model for training to obtain the teacher model weights;
[0013] S3. Input the preprocessed training set into the student model and the trained teacher model respectively, and initialize the pixel memory and region memory.
[0014] S4. After adapting to the heterogeneous architecture, the distillation strategy is iteratively trained through a multi-level knowledge transfer module. The five-dimensional loss weighted fusion system includes: foreground-sensitive segmentation loss, output layer knowledge distillation loss, intermediate layer knowledge distillation loss, mini-batch pixel comparison loss, and memory bank comparison loss.
[0015] S5. After the distillation strategy training is completed, the preprocessed validation set is used to evaluate the average intersection-union ratio mIoU, average Dice coefficient mDice, precision, and recall of the student model after distillation. After the validation is successful, the student model is deployed in the terminal that is compatible with the acquisition device.
[0016] S6. After the acquisition device acquires the lesion image in real time, the terminal completes pixel-level segmentation of the lesion image through the embedded student model and outputs the corresponding medical image segmentation result.
[0017] Preferably, in S1, the teacher model adopts the deep learning model HRNet semantic segmentation model with HRNetv2-W48 as the backbone network, and the student model adopts the image segmentation model ResNet50-UNet.
[0018] The training parameters for the distillation strategy include: using the Adam optimizer, setting the initial learning rate to 0.001, the weight decay to 1e-5, the training batch size to 6, the training epochs to 100, and using cosine annealing learning rate scheduling.
[0019] Preferably, the preprocessing of the training set and validation set involves uniformly scaling the images to 480×480, with the training set employing anatomical structure preservation and enhancement, and the validation set only being scaled and normalized.
[0020] Preferably, in S4, the foreground-sensitive segmentation loss is used. It is characterized by the following formula:
[0021]
[0022] In the above formula, DiceLoss represents the overlap between the predicted segment and the real label, FocalLoss is the target detection design, α is the class weight that emphasizes the lesion pixel, and γ is the focusing parameter applied to the focal loss.
[0023] Preferably, in S4, the output layer knowledge distillation loss is calculated on the output KL divergence loss of the teacher model and the student model, and then the output distribution of the teacher model and the student model is aligned based on the KL divergence, so that the student model can learn the judgment logic of the teacher's high-level semantics.
[0024] Among them, the output layer knowledge distillation loss It is characterized by the following formula:
[0025]
[0026] In the above formula, τ is the temperature parameter, and S T For the output of the teacher model, S S The output of the student model is represented by , softmax is the activation function that converts the numerical values of the model output into a probability distribution, and KL() is the divergence used to measure the difference between the distributions of the teacher and student models. This represents a separator.
[0027] Preferably, in S4, the intermediate layer knowledge distillation loss constructs an adaptive projection matrix through a dynamic dimension mapping algorithm, and aligns the heterogeneous feature spaces of the teacher model and the student model based on a core mechanism of key-value caching mechanism + dynamic convolution construction + instant interpolation space alignment, thereby conveying the local structural details of the teacher model at different levels.
[0028] The key-value caching mechanism is based on a combination of feature dimension and size parameter (C). in C out H in W in H out W out The key is used to query the preset projector cache dictionary; if the corresponding key exists in the cache dictionary, the projector is reused directly; otherwise, the dynamic convolution module construction process is triggered. in H represents the number of feature input channels in the teacher model. in W in C represents the input space size of the teacher model. out H represents the number of feature target channels in the student model. out W out The target space size for the student model;
[0029] The dynamic convolution construction includes: when C out >C in At that time, a concatenated structure of 3×3 convolution aggregation + 1×1 convolution mapping is adopted, where the number of intermediate channels C mid =min(C out C in ×4), and limit the channel expansion limit to 4 times the number of input channels;
[0030] When C out ≤C in At that time, 1×1 convolution is used to directly complete channel projection, all convolutional layers are configured with BatchNorm2d for normalization, and no bias term is set;
[0031] The instantaneous interpolation space alignment refers to: after completing the channel dimension mapping through dynamic convolution, if the projected feature space size and the target size (H) are aligned, then the instantaneous interpolation space alignment is achieved. out W out If they are inconsistent, bilinear interpolation is used to adjust the size in real time.
[0032] Preferably, in S4, the intermediate layer knowledge distillation loss It is characterized by the following formula:
[0033]
[0034] in, , , , are the features of the i-th and i+1-th layers of the teacher model and student model, respectively; proj is the heterogeneous feature projection operation based on the adaptive projection matrix; and Loss() represents the feature matching loss function.
[0035] Preferably, in S4, mini-batch pixel contrast loss It is characterized by the following formula:
[0036]
[0037] In the above formula, p i p j These are the pixel-level feature vectors corresponding to pixel i and pixel j, respectively, sin(p i ,p j ) represents the similarity between pixel i and pixel j, k represents the index of all pixels in the mini-batch, τ is the temperature parameter, and exp() represents the natural exponential function.
[0038] Preferably, in S4, during the distillation strategy training, hard samples are identified based on the dual threshold method, and a dynamic weight allocation mechanism is adopted to dynamically adjust the sampling frequency of the hard sample region as the training rounds change.
[0039] The dual thresholds include: a prediction probability threshold and a pixel error threshold.
[0040] The present invention has at least the following beneficial effects: The knowledge distillation strategy of the present invention can realize the transfer of anatomical prior knowledge from complex teacher models to lightweight student models, which is suitable for resource-constrained clinical deployment scenarios. It should be noted that resource constraints here refer to the limitation of equipment configuration, that is, the existing technology has the problem of insufficient equipment computing power. For example, most teacher models usually have a large number of parameters, which makes the equipment requirements for teacher models higher in practical applications. However, even on high-configuration equipment, the processing efficiency is low. The student model after distillation improves the accuracy based on the existing student model, but still maintains a low number of parameters. Therefore, it can still ensure both accuracy and efficiency on equipment with general configuration. Its effects can be concretely described as including:
[0041] Firstly, this invention proposes a dynamic dimension mapping algorithm that automatically identifies the feature dimension differences between the teacher model and the student model, constructs an adaptive projection matrix, and accurately aligns the intermediate layer feature space. This design does not rely on feature map size consistency and achieves "heterogeneous architecture feature transfer" for the first time in the field of medical image segmentation, filling the gap in existing technologies and published patents in heterogeneous model distillation and adaptation technology, and providing a fundamental guarantee for multi-level knowledge transfer.
[0042] Secondly, this invention constructs a dual-database mechanism of pixel memory (capacity 2000) + region memory (capacity 800). By comparing the memory database loss (pixel level + region level), the student model is forced to reuse the teacher's historical structural knowledge throughout the training process, solving the core problem of "knowledge fragmentation" in existing technologies and making the student model more in line with clinical anatomy logic.
[0043] Third, this invention designs a five-dimensional loss-weighted fusion system. Compared with the single-dimensional design of existing technologies, this system achieves "full-scenario coverage" of knowledge transfer, enabling the student model to be upgraded from "partial learning" to "comprehensive ability transfer".
[0044] Fourth, this invention proposes a dual-threshold hard sample identification + dynamic weight allocation strategy to accurately locate key clinical hard sample regions while avoiding overfitting.
[0045] Fifth, based on the full transfer of knowledge, this invention achieves a clinically adapted design of "architectural optimization + efficiency quantification". It can retain the single-path encoding-decoding structure of the student model (UNet) and avoid the memory access redundancy of multi-branch models, ensuring memory friendliness. Through a multi-level distillation strategy, while improving accuracy, the inference speed of the student model (UNet) is increased by 2.3 times compared with the teacher model (HRNet), clearly meeting the real-time requirement of ≥20 FPS for mobile devices. Moreover, it maintains high segmentation accuracy and has better generalization ability than existing technologies in complex clinical scenarios such as small sample size (BIOHit has only 300 images), low contrast (tongue images), and blurred boundaries (skin texture boundaries).
[0046] Other advantages, objectives and features of the present invention will become apparent in part from the following description, and in part from those skilled in the art through study and practice of the invention. Attached Figure Description
[0047] Figure 1 This is a system block diagram of the medical image segmentation method based on the heterogeneous distillation strategy of the present invention. Detailed Implementation
[0048] The present invention will now be described in further detail with reference to the accompanying drawings, so that those skilled in the art can implement it based on the description.
[0049] Based on the full transfer of knowledge, this invention designs a distillation strategy to achieve a clinically adapted design of "architecture optimization + efficiency quantification". At the architecture level, the distillation strategy retains the single-path encoder-decoder structure of the student model (UNet) to avoid memory access redundancy in multi-branch models and ensure memory friendliness. At the efficiency level, through a multi-level distillation strategy, while improving accuracy, the inference speed of the student model (UNet) is improved by approximately 12.5 times compared to the teacher model (HRNet) (185.99 FPS vs. 14.92 FPS, measured on RTX-4090D, PyTorch 2.1), clearly meeting the real-time requirement of ≥20 FPS for mobile devices. At the robustness level, it maintains high segmentation accuracy and has better generalization ability than existing technologies even in complex clinical scenarios such as small sample sizes (BIOHit has only 300 images), low contrast (tongue images), and blurred boundaries (skin texture boundaries).
[0050] Specifically, such as Figure 1 As shown, this invention proposes a heterogeneous architecture memory-enhanced knowledge distillation strategy (hereinafter referred to as: distillation strategy or HAMKD). Through three core modules—a heterogeneous architecture adaptation module, a memory enhancement module, and a multi-level knowledge transfer module—it achieves efficient knowledge transfer from the teacher model to the student model. The specific scheme is as follows:
[0051] 1. The heterogeneous architecture adaptation module automatically identifies the feature dimension differences between the teacher model (multi-resolution branch) and the student model (single path encoding-decoding) by designing a dynamic dimension mapping algorithm, constructs an adaptive projection matrix, aligns the intermediate layer feature space, and ensures that the features of the heterogeneous architecture can be transferred, achieving a breakthrough from "same architecture dependency" to "dynamic compatibility".
[0052] 2. The memory enhancement module mainly constructs a pixel memory bank (size 2000) and a region memory bank (size 800) to store the historical pixel features (such as lesion boundary pixels) and semantic region features (such as lesions and background regions) of the teacher model, respectively, to achieve knowledge storage and reuse, that is, from "single matching" to "continuous transmission", providing continuous structural knowledge reference for the student model.
[0053] 3. Multi-level knowledge transfer: By designing multi-dimensional loss functions, including segmentation loss, output layer distillation loss, intermediate layer distillation loss, mini-batch pixel contrast loss, and memory bank contrast loss (pixel level + region level), a five-dimensional loss weighted fusion system is constructed, which forces the student model to learn teacher knowledge from multiple dimensions such as output, intermediate layer, pixel, and region.
[0054] Furthermore, the architecture and parameter settings of the HAMKD distillation strategy of this invention mainly include:
[0055] 1. Regarding model selection, the teacher model is HRNet based on HRNetv2-W48, and the student model is UNet based on ResNet50;
[0056] 2. In terms of training parameter settings, the Adam optimizer (initial learning rate 0.001, weight decay 1e-5), batch size 6, training epochs 100, and cosine annealing learning rate scheduling are used.
[0057] 3. Regarding loss parameter settings, a five-dimensional loss weighted fusion system is designed. Compared to the single-dimensional design of existing technologies, this system achieves "full-scenario coverage" of knowledge transfer, upgrading the student model from "partial learning" to "comprehensive ability transfer." Specifically, the five-dimensional loss weighted fusion system includes: segmentation loss. Output layer distillation loss Intermediate layer distillation loss Small batch pixel contrast loss Memory contrast loss ( , The total loss formula is represented by the following equation:
[0058]
[0059] In the above formula, , , , , , They are respectively , , , , , Each has its corresponding weight.
[0060] 1) Segmentation loss
[0061] To address the severe class imbalance problem in medical image segmentation, a foreground-sensitive segmentation loss function is designed:
[0062]
[0063] This approach integrates Dice and Focal losses. Weighted Dice Loss measures the overlap between the predicted segmentation and the ground truth label (focusing on region integrity, suitable for morphological matching of lesion regions in medical images). Focal Loss is designed to focus on target detection of hard-to-classify pixels. Its core principle is to reduce the weight of easily classified samples and increase the attention given to hard-to-classify samples, thus alleviating class imbalance (low proportion of foreground / lesion pixels in medical images) and forcing the model to prioritize learning foreground regions. In practical applications, this step addresses the class imbalance problem in medical images where the foreground (lesion) proportion is low. Dice Loss ensures region integrity, while Focal Loss focuses on hard-to-classify pixels (such as blurred boundaries). The weighting scheme is designed as follows: DiceLoss weight 0.6 + FocalLoss weight 0.4. These weights are determined based on the characteristics of medical data, specifically through analysis of the foreground pixel proportions in five medical datasets, and are used to balance the learning of "region integrity" and "hard-to-classify pixels." The category weight α is used to emphasize lesion pixels. A focusing parameter γ=2.0 is applied to the focus loss to adaptively reduce the pixel weights.
[0064] 2) Knowledge distillation loss in the output layer
[0065] Calculate teacher output S T With students S S KL divergence loss
[0066]
[0067] In the above formula, τ is the temperature parameter, and S T For the output of the teacher model, S S For the output of the student model, The delimiter is `softmax()`, which is the activation function that converts the model's output values into a probability distribution. `KL()` represents the KL divergence, used to measure the difference between the teacher and student model distributions, serving as a loss to guide the student model in learning the teacher model's knowledge. Based on the KL divergence, the teacher and student output distributions are aligned, with a temperature parameter τ=1.5 to soften the probability distribution and enhance knowledge transfer (the lower the temperature, the "harder" the distribution, focusing on high-confidence predictions). The model smooths the teacher's output using the temperature τ. T With student S S The output distribution of the model aligns the output probability distribution of the student model with that of the teacher, enabling the student to learn the teacher's judgment logic regarding high-level semantics (such as lesion category and overall structure).
[0068] 3) Knowledge distillation loss in the intermediate layer
[0069] An adaptive projection matrix is constructed using a dynamic dimension mapping algorithm. Based on a three-in-one core mechanism of "key-value caching + dynamic convolution + instant interpolation", the heterogeneous feature spaces of teacher and student models are aligned, conveying the local structural details of the teacher model at different levels (Stage 3 focuses on texture features, and Stage 4 focuses on morphological features). Specifically, the feature losses of Stage 3 and Stage 4 are calculated using the following formula:
[0070]
[0071] in, , , , These are the features of the teacher (HRNet, multi-resolution parallel branch) and the student (UNet, encoder-decoder symmetric structure) at layers i and i+1, respectively. Loss() represents the feature matching loss function, and proj is the heterogeneous feature projection operation based on the adaptive projection matrix. The specific implementation process is as follows:
[0072] Feature Dimension and Size Recognition: Extract Stage 3 and Stage 4 features from teachers and students, and automatically parse the number of teacher feature input channels C. in Input space dimensions (H) in W in ), and the number of student characteristic target channels C out Target space size (H) out W out ( ), clearly define the differences in dimensions and sizes.
[0073] Key-value caching mechanism adaptation: based on a combination of feature dimension and size parameters (C in C out H in W in H out W out The key is used to query the preset projector cache dictionary; if the corresponding projector exists in the cache, it is reused directly; otherwise, the dynamic convolution module construction process is triggered.
[0074] Construction of dynamic convolution module: based on C in With C out The size relationship is used to dynamically generate the corresponding convolutional structure: when C out >C in At that time, a concatenated structure of "3×3 convolution first aggregates local context + 1×1 convolution projection" is adopted, where the number of intermediate channels C mid =min(C out C in ×4), limiting the channel expansion limit to 4 times the number of input channels; when C out ≤C inIn this case, 1×1 convolution is used to directly complete the channel mapping. All convolutional layers are configured with BatchNorm2d for normalization and no bias term is set.
[0075] Instantaneous interpolation space alignment: After completing channel dimension mapping through dynamic convolution, if the size of the projected feature space is similar to the target size (H... out W out If there is a discrepancy, bilinear interpolation is used to adjust the size in real time.
[0076] Joint parameter optimization: The convolutional layer weights of the adaptive projection matrix are trained synchronously with the student model parameters. The difference between the projected features and the student features is calculated based on "SSIM loss + L1 loss" (weight ratio 0.7:0.3). The convolutional layer weights are updated through backpropagation to continuously optimize the mapping accuracy of heterogeneous features.
[0077] 4) Small batch pixel contrast loss
[0078] This step calculates intra-batch pixel associations using sinusoidal similarity, maximizing the similarity between pixels of the same type (e.g., lesion boundary pixels) and minimizing the similarity between pixels of different types (e.g., lesion and background pixels), thereby improving fine-grained segmentation accuracy and conveying detailed knowledge of local structures. Specifically, the mini-batch pixel comparison calculates the intra-batch pixel similarity sin(p i ,p j Therefore, the formula for small-batch pixel contrast loss can be characterized as follows:
[0079]
[0080] In the above formula, p i p j These are the pixel-level feature vectors corresponding to pixel i and pixel j, respectively, sin(p i ,p j Let represent the similarity between pixel i and pixel j, calculated using a sine function. exp() represents the natural exponential function, k represents the index of all pixels in the mini-batch (i.e., any pixel within the mini-batch), and τ represents the temperature. A pixel similarity comparison is constructed based on a temperature τ=0.25. The overall formula maximizes the similarity of pixels of the same type and minimizes the similarity of pixels of different types, allowing students to learn the teacher's discrimination logic for fine-grained pixels (such as boundaries), thus improving the segmentation accuracy of fine structures such as lesion edges.
[0081] 5) Memory comparison
[0082] This stage primarily involves constructing a dual-database mechanism of pixel memory (capacity 2000) and region memory (capacity 800) to reuse historical knowledge and enhance the stability of structural feature learning. The pixel memory stores key pixel features such as lesion boundaries and subtle textures from the teacher model's historical data, providing continuous boundary discrimination knowledge references for the student model and preventing the "learning and forgetting of boundary features." The region memory stores semantic region features such as lesions, background, and normal tissue, conveying spatial correlation knowledge such as the relative position of lesions and surrounding tissues, and the proportion of regional morphology, thus strengthening the reuse of anatomical priors. Through contrastive loss (pixel-level + region-level), the student model is forced to reuse the teacher's historical structural knowledge throughout training, addressing the core problem of "knowledge fragmentation" in existing technologies and making the student model more aligned with clinical anatomical logic.
[0083] Pixel level The system utilizes a pixel memory bank (size 2000) to store historical pixel features, transmits boundary discrimination knowledge, constrains the consistency between student pixel features and teacher pixel features, and strengthens the learning of key pixels such as lesion boundaries.
[0084] regional level Based on a region memory (size 800), it divides semantic regions and stores semantic region features (such as lesions and background), and transmits spatial relationship knowledge between regions (such as the relative position of lesions and normal tissues), thereby improving the integrity of region segmentation.
[0085] Furthermore, it should be noted that traditional distillation employs an "equal sampling" strategy, assigning the same learning weight to all samples (including clear, easily segmented regions and blurry, difficult-to-segment regions). This causes the model to focus more on easily segmented regions, resulting in insufficient learning of difficult-to-segment small lesions and low-contrast regions (such as the edges of the tongue). Recent structure-aware distillation methods do not have a dedicated optimization mechanism for hard samples, only indirectly improving the situation by increasing the structural loss weight, which cannot accurately locate hard sample regions and has low optimization efficiency. To address this issue, this invention employs a hard sample oversampling method during training, namely, identifying hard samples based on a dual threshold method (prediction probability threshold + pixel error threshold). For example, in skin lesion segmentation, when the probability of a region being predicted as a lesion is less than 0.6 and the pixel error with the true label is greater than 0.3, it is determined to be a hard sample (such as blurry lesion boundaries or small lesion regions). Furthermore, a "dynamic weight allocation" mechanism is adopted, which dynamically optimizes the sampling weight of hard samples according to the training stage: for example, a weight of 1.5 times is set in rounds 1-30 (early stage) to prioritize the establishment of basic knowledge of hard samples; a weight of 2 times is set in rounds 31-70 (mid stage) to strengthen the learning of details of hard samples; and a weight of 1.2 times is set in rounds 71-100 (late stage) to avoid overfitting.
[0086] Example:
[0087] S1. Data Preprocessing and Partitioning
[0088] Multiple medical datasets (ISIC-2016 / 2017, BIOHit) were selected and divided into training and validation sets by stratified sampling at a ratio of 9:1. Images were uniformly scaled to 480×480. The training set was enhanced with "anatomical structure preservation", while the validation set was only scaled and normalized.
[0089] S2. Teacher model pre-training: Train the HRNet teacher model by inputting preprocessed training set images, training for 100 rounds, and saving the weights and memory initialization parameters of the trained teacher model.
[0090] The preprocessed training set images are simultaneously input into the pre-trained HRNet teacher model and the ResNet50-based UNet student model, and pixel memory and region memory with capacities of 2000 and 800 respectively are initialized to clear historical data.
[0091] S3, HAMKD distillation training
[0092] S31 enters heterogeneous architecture adaptation phase
[0093] S310. Automatically extract the intermediate layer features of Stage3 and Stage4 of the teacher-student model and analyze the differences between them in terms of channel number, spatial size, etc.
[0094] S311. Use the combination of feature parameters as the key to query the projector cache dictionary. If the corresponding projector exists, it is reused directly. If it does not exist, the convolution structure is dynamically constructed according to the relationship between the number of channels (when the target number of channels is greater than the input value, a "3×3 convolution + 1×1 convolution" concatenated structure is used, and when the target value is less than or equal to the input value, a 1×1 convolution is used).
[0095] S312. By adjusting the feature space size through bilinear interpolation, accurate alignment of the heterogeneous feature spaces of teachers and students can be achieved.
[0096] S313. Perform knowledge storage update. Extract key pixel features such as lesion boundaries and subtle textures from the output features of the teacher model and store them in the pixel memory bank. At the same time, perform semantic region division of the teacher features and extract global features of each region and store them in the region memory bank to provide continuous reference for subsequent knowledge transfer.
[0097] S32. Initiating Multi-Level Knowledge Transfer
[0098] S320. Calculate the segmentation loss consisting of DiceLoss and FocalLoss with a weight of 0.6:0.4 to balance region integrity with learning of hard-to-classify pixels;
[0099] S321. Based on temperature parameter 1.5, calculate the output layer distillation loss using KL divergence to align with the high-level semantic judgment of teachers and students;
[0100] S322. Calculate the intermediate layer distillation loss of Stage 3 and Stage 4 by combining the heterogeneous feature projection results to convey local structural details;
[0101] S323. Calculate the small batch pixel contrast loss using pixel similarity with a temperature parameter of 0.25 to improve fine-grained segmentation accuracy;
[0102] S324. Calculate the pixel-level and region-level memory contrast loss to constrain the consistency between student features and teacher historical features, and use the following formula to fuse these five-dimensional losses according to preset weights into a total loss, so as to synchronously update the student model parameters and adaptive projection matrix weights through backpropagation:
[0103]
[0104] During training, hard samples are optimized by oversampling hard samples, achieving a breakthrough from "equal neglect" to "precise enhancement". That is, hard samples are identified based on the dual threshold method (prediction probability threshold + pixel error threshold). For example, in skin lesion segmentation, when the probability of a certain area being predicted as a lesion is less than 0.6 and the pixel error with the real label is greater than 0.3, it is judged as a hard sample area such as blurred lesion boundary or micro lesion area.
[0105] Meanwhile, sampling weights are dynamically allocated according to the training rounds. That is, the sampling frequency of hard sample regions is dynamically adjusted with the training rounds. In the initial rounds (rounds 1-30), the sampling weight is 1.5 times that of ordinary samples, in the middle rounds (rounds 31-70), it is increased to 2 times, and in the later rounds (rounds 71-100), it drops back to 1.2 times, thereby strengthening hard sample learning and avoiding overfitting.
[0106] This example demonstrates significant improvements in the student model (UNet) across multiple datasets through multi-level knowledge transfer and memory enhancement. As shown in Table 1, on the BIOHit dataset, the student model's average intersection-over-union (mIoU) improved from 93.38% to 98.02%, and the average Dice coefficient (mDice) improved from 96.54% to 99.00%.
[0107] Table 1: Performance comparison of the model before and after knowledge distillation on the BIOHit dataset.
[0108]
[0109] As shown in Table 2, on the ISIC-2016 dataset, the student model's mIoU improved from 70.21% to 87.78%, and mDice improved from 81.33% to 93.41%; on the ISIC-2017 dataset, the student model's mIoU improved from 68.70% to 77.93%, and mDice improved from 80.70% to 86.95%.
[0110] Table 2: Comparative analysis of model performance before and after knowledge distillation on the ISIC-2016 and ISIC-2017 datasets.
[0111] S4. Student Model Validation and Deployment
[0112] The mean intersection-union ratio (mIoU), mean Dice coefficient (mDice), precision, and recall of the student model after distillation are evaluated using a validation set. After successful validation, the student model is deployed on a resource-constrained mobile device to test the inference speed (≥20 FPS).
[0113] This example uses a student model distilled from HAMKD, retaining the single-path structure of UNet, which is memory-friendly and offers a 2.3x speed improvement in inference compared to the original HRNet, making it deployable on resource-constrained mobile devices. Furthermore, in scenarios with small samples (BIOHit, 300 images), low contrast (tongue images), and blurred boundaries (skin texture boundaries), the distilled student model maintains high robustness and demonstrates superior generalization ability compared to traditional distillation methods. Further, taking tongue image detection as an example, the heterogeneous distillation-based high-precision transfer and lightweight method involved in this invention is directly applied to the medical segmentation scenario assisting in intelligent tongue diagnosis in Traditional Chinese Medicine. Specifically, a high-precision heavy heterogeneous network is selected as the teacher model (this high-precision heavy heterogeneous network possesses accurate pixel-level segmentation capabilities and can fully learn the edge features of the tongue, the boundary features between the tongue coating and the tongue body, and the subtle texture features of tongue lesions, etc., which are core pathological knowledge of tongue diagnosis); a lightweight medical segmentation network is used as the heterogeneous student model. Through heterogeneous distillation, the high-precision segmentation knowledge of the teacher model is completely transferred to the student model, while network pruning and weight adjustments are performed on the student model. Lightweight optimizations, such as quantization, compress the number of model parameters while ensuring segmentation accuracy. The compressed student model can be directly deployed in the processing terminal corresponding to portable / mobile image acquisition devices. During clinical tongue examination, the patient extends their tongue, and after the tongue image is acquired in real time by the camera, the lightweight model can perform rapid inference calculations in the student model on the local processing terminal, performing higher-precision pixel-level segmentation on the acquired tongue image: accurately identifying the position and contour of the tongue, clearly distinguishing the junction area of the tongue with the oral cavity and face, effectively detecting the tongue, eliminating background interference, and the segmentation inference time of a single tongue image is ≤50ms, meeting the real-time hard requirement of ≥20 FPS. This realizes the clinical closed-loop application of tongue image segmentation acquisition and local real-time high-precision segmentation, greatly improving the accuracy of subsequent intelligent diagnosis.
[0114] The above solution is merely an illustration of a preferred example and is not limited thereto. When implementing this invention, appropriate substitutions and / or modifications can be made according to the user's needs.
[0115] Although embodiments of the present invention have been disclosed above, they are not limited to the applications listed in the specification and embodiments. It can be applied to various fields suitable for the present invention. Other modifications can be readily made by those skilled in the art. Therefore, without departing from the general concept defined by the claims and their equivalents, the present invention is not limited to the specific details and examples shown and described herein.
Claims
1. A medical image segmentation method based on a heterogeneous distillation strategy, characterized in that, include: S1. Construct a distillation strategy that complements the teacher model and the student model. The distillation strategy includes: a heterogeneous architecture adaptation module based on a dynamic dimension mapping algorithm, a memory enhancement module based on pixel memory bank and region memory bank, and a multi-level knowledge transfer module based on a five-dimensional loss weighted fusion system. S2. Input the preprocessed training set into the teacher model for training to obtain the teacher model weights; S3. Input the preprocessed training set into the student model and the trained teacher model respectively, and initialize the pixel memory and region memory. S4. After adapting to the heterogeneous architecture, the distillation strategy is iteratively trained through a multi-level knowledge transfer module. The five-dimensional loss weighted fusion system includes: foreground-sensitive segmentation loss, output layer knowledge distillation loss, intermediate layer knowledge distillation loss, mini-batch pixel comparison loss, and memory bank comparison loss. S5. After the distillation strategy training is completed, the preprocessed validation set is used to evaluate the average intersection-union ratio mIoU, average Dice coefficient mDice, precision, and recall of the student model after distillation. After the validation is successful, the student model is deployed in the terminal that is compatible with the acquisition device. S6. After the acquisition device acquires the lesion image in real time, the terminal completes the pixel-level segmentation of the lesion image through the embedded student model and outputs the corresponding medical image segmentation result. In S4, the intermediate layer knowledge distillation loss constructs an adaptive projection matrix through a dynamic dimension mapping algorithm, and aligns the heterogeneous feature spaces of the teacher model and the student model based on a core mechanism of key-value caching mechanism + dynamic convolution construction + instant interpolation space alignment, thereby conveying the local structural details of the teacher model at different levels. The key-value caching mechanism is based on a combination of feature dimension and size parameter (C). in C out H in W in H out W out The key is used to query the preset projector cache dictionary; if the corresponding key exists in the cache dictionary, the projector is reused directly; otherwise, the dynamic convolution module construction process is triggered. in H represents the number of feature input channels in the teacher model. in W in C represents the input space size of the teacher model. out H represents the number of feature target channels in the student model. out W out The target space size for the student model; The dynamic convolution construction includes: when C out >C in At that time, a concatenated structure of 3×3 convolution aggregation + 1×1 convolution mapping is adopted, where the number of intermediate channels C mid =min(C out C in ×4), and limit the channel expansion limit to 4 times the number of input channels; When C out ≤C in At that time, 1×1 convolution is used to directly complete channel projection, all convolutional layers are configured with BatchNorm2d for normalization, and no bias term is set; The instantaneous interpolation space alignment refers to: after completing the channel dimension mapping through dynamic convolution, if the projected feature space size and the target size (H) are aligned, then the instantaneous interpolation space alignment is achieved. out W out If they are inconsistent, bilinear interpolation is used to adjust the size in real time.
2. The medical image segmentation method based on a heterogeneous distillation strategy as described in claim 1, characterized in that, In S1, the teacher model uses the deep learning model HRNet semantic segmentation model with HRNetv2-W48 as the backbone network, and the student model uses the image segmentation model ResNet50-UNe. The training parameters for the distillation strategy include: using the Adam optimizer, setting the initial learning rate to 0.001, the weight decay to 1e-5, the training batch size to 6, the training epochs to 100, and using cosine annealing learning rate scheduling.
3. The medical image segmentation method based on a heterogeneous distillation strategy as described in claim 1, characterized in that, The preprocessing of the training and validation sets involves uniformly scaling the images to 480×480, with the training set using anatomical structure preservation and enhancement, and the validation set only being scaled and normalized.
4. The medical image segmentation method based on a heterogeneous distillation strategy as described in claim 1, characterized in that, In S4, foreground-sensitive segmentation loss It is characterized by the following formula: In the above formula, DiceLoss represents the overlap between the predicted segment and the real label, FocalLoss is the target detection design, α is the class weight that emphasizes the lesion pixel, and γ is the focusing parameter applied to the focal loss.
5. The medical image segmentation method based on a heterogeneous distillation strategy as described in claim 1, characterized in that, In S4, the output layer knowledge distillation loss is calculated by aligning the output distributions of the teacher model and the student model based on the KL divergence loss of the teacher model and the student model, so that the student model can learn the judgment logic of the teacher's high-level semantics. Among them, the output layer knowledge distillation loss It is characterized by the following formula: In the above formula, τ is the temperature parameter, and S T For the output of the teacher model, S S The output of the student model is represented by `softmax()`, which is the activation function that converts the numerical values of the model output into a probability distribution. `KL()` is the divergence used to measure the difference between the teacher and student model distributions. This represents a separator.
6. The medical image segmentation method based on a heterogeneous distillation strategy as described in claim 1, characterized in that, In S4, intermediate layer knowledge distillation loss It is characterized by the following formula: in, , , , are the features of the i-th and i+1-th layers of the teacher model and student model, respectively; proj is the heterogeneous feature projection operation based on the adaptive projection matrix; and Loss() represents the feature matching loss function.
7. The medical image segmentation method based on a heterogeneous distillation strategy as described in claim 1, characterized in that, In S4, mini-batch pixel contrast loss It is characterized by the following formula: In the above formula, p i p j These are the pixel-level feature vectors corresponding to pixel i and pixel j, respectively, sin(p i ,p j ) represents the similarity between pixel i and pixel j, k represents the index of all pixels in the mini-batch, τ is the temperature parameter, and exp() represents the natural exponential function.
8. The medical image segmentation method based on a heterogeneous distillation strategy as described in claim 1, characterized in that, In S4, the memory-based contrast loss includes: pixel-level loss. Regional losses ; Among them, pixel-level loss Based on the historical pixel features stored in the pixel memory bank, boundary discrimination knowledge is transmitted, thereby constraining the consistency between student pixel features and teacher pixel features to strengthen the learning of key pixels. Regional losses By storing semantic region features in the segmented semantic regions of the region memory, the spatial relationship knowledge between regions can be transmitted to improve the integrity of region segmentation.
9. The medical image segmentation method based on a heterogeneous distillation strategy as described in claim 1, characterized in that, In S4, during the distillation strategy training, hard samples are identified based on the dual threshold method, and a dynamic weight allocation mechanism is adopted to dynamically adjust the sampling frequency of hard sample regions as the training rounds change. The dual thresholds include: a prediction probability threshold and a pixel error threshold.
Citation Information
Patent Citations
Medical image segmentation method, system and equipment based on knowledge distillation
CN117274282A
Heterogeneous model alignment method and system based on cross-modal knowledge distillation
CN120724393A