Relational knowledge distillation method based on category boundary uncertainty

Abstract

The invention discloses a relational knowledge distillation method based on category boundary uncertainty. The method comprises the following steps of: 1) simultaneously and respectively processing an image through feature encoders of two network branches to obtain an intermediate feature map, and processing the intermediate feature map through a decoder to obtain a final segmentation prediction map; 2) calculating the consistency distillation loss of the characteristic level based on the intermediate characteristic spectrum obtained in the step 1); 3) calculating pixel-level consistency distillation loss based on the segmentation prediction map obtained in the step 1); 4) respectively calculating the target loss of the branch 1 and the target loss of the branch 2 based on the double-branch segmentation prediction map obtained in the step 1) and the corresponding real category label A; 5) calculating a total optimization loss function; and 6) training the model based on the loss function in the step 5). According to the method, the BURKD algorithm is introduced, the relationship between samples and between categories can be better modeled, adaptive dynamic weighting is carried out on each category by adopting a method based on category boundary uncertainty, and the semantic segmentation performance is improved.

Description

Technical Field

[0001] The present invention belongs to the field of relational knowledge distillation, and specifically constructs a relational knowledge distillation method based on category boundary uncertainty to assist in enhancing a semantic segmentation method. Background Art

[0002] Knowledge Distillation (KD) aims to transfer valuable knowledge from a large teacher model to a small student model to improve the inference accuracy of the student model. Generally speaking, distillation methods can be divided into three types: response-based distillation (matching teacher and student soft scores), representation-based distillation (matching feature maps), and relation-based distillation (matching sample relations). The key factors for the success of KD are: 1) the type of knowledge distillation (e.g., feature or score level), 2) the adopted distillation strategy (e.g., offline distillation, online distillation, and self-distillation), and 3) the architecture of the teacher and student (e.g., multi-teacher KD, attention-based KD, and graph-based KD).

[0003] Although knowledge distillation originated from resource-efficient deep learning, it has been more widely used in fields such as incomplete multimodal learning. For incomplete multimodal medical image segmentation tasks, the model is highly sensitive to feature representation, and the rough and general KD may not be effective enough and can be considered purposeless and redundant, with a large performance gap between the obtained results and the expected results. In addition, due to the different inputs of the teacher and student networks, transferring knowledge through representation-based methods may lead to overfitting. Recent methods focus on transferring specific modality information through relationship-based methods. However, current methods usually consider different instances equally and ignore their specificity, which will lead to suboptimal performance.

[0004] In view of the above limitations, we proposed a relational knowledge distillation method based on category boundary uncertainty. We introduced this method into the semantic segmentation network and verified the effectiveness of this method in the incomplete multimodal brain tumor segmentation scenario. Summary of the invention

[0005] In order to solve the problem in the background technology, the present invention provides a relational knowledge distillation method based on category boundary uncertainty. The purpose of the present invention is to judge the importance of instance categories according to entropy through a relational knowledge distillation method based on category boundary uncertainty, thereby giving different weights to different instances, so as to specifically consider different semantic segmentation instance categories, thereby enhancing the accuracy of semantic segmentation.

[0006] The technical solution adopted by the present invention is as follows, which specifically comprises the following steps:

[0007] 1) Input an image into two network branches at the same time, first pass through the feature encoder Get the intermediate feature maps F1 and F2; then pass through the decoder Get the final segmentation prediction map

[0008] 2) Based on the intermediate feature maps F1 and F2 obtained in step 1), calculate the feature-level consistency distillation loss L fc ;

[0009] 3) Based on the segmentation prediction map obtained in step 1) Compute pixel-level consistency distillation loss

[0010] 4) Based on the dual-branch segmentation prediction map obtained in step 1) and its corresponding true category label A, and calculate the target loss of the first network branch and the second network branch respectively

[0011] 5) Combine the L obtained in step 2) fc and step 3) And the result obtained in step 5) Calculate the total optimization loss L i ;

[0012] 6) Train the model based on the loss function of step 5);

[0013] 7) Use the model trained in step 6) to obtain the segmentation or classification results of the image to be processed.

[0014] In the step 1):

[0015] The first network branch includes a feature encoder Decoder

[0016] The second network branch includes a feature encoder Decoder

[0017] The step 2) is specifically as follows: based on the intermediate feature maps F1 and F2 of the two branches obtained in step 1), the category-specific prototypes, the contextual relationships between modeled samples and categories, and the re-weighting of the uncertainty of category boundaries are calculated in sequence to obtain the feature-level consistency loss L fc , the specific operations are as follows:

[0018] Step 2.1) Category-specific prototype generation: Given an intermediate feature map And its corresponding true label A∈RH×W×Z , calculate the class-specific prototypes of the first network branch and the second network branch respectively:

[0019]

[0020] in, H i ,W i ,Z i Indicates the length, width, and height of the feature, d f Indicates the number of channels of the feature, H, W, Z represent the length, width, and height of the real label A, N c represents the number of pixels belonging to category c, Is an indicator function, which means that the value is 1 when the internal condition is True, otherwise it is 0;

[0021] Step 2.2) Relationship context modeling between samples and categories: Based on the category-specific prototypes obtained in step 2.1), the relationship maps R1 and R2 between samples and categories are obtained by the following formula:

[0022]

[0023] in, Respectively represent the prototype of the c1th category of the b1th sample in the first network branch and the second network branch, R1, R2∈R (B*C)×(B*C) , B represents the number of samples in a batch, and C represents the number of semantic categories;<a,b> Indicates the dot product operation of computing elements a and b, R i (b1, c1; b2, c2) represents the feature similarity between the c1 category of the b1th sample and the c2 category of the b2th sample in the i-th network branch; then the difference in the relationship distance between samples and categories between the two branches is calculated:

[0024]

[0025] in, represents the overall relationship difference of the b1th sample and c1th category between the first network branch and the second network branch;

[0026] Step 2.3) Re-weighting based on boundary uncertainty: In order to prioritize consistency learning on more challenging categories, we introduce a category boundary-specific uncertainty re-weighting method; based on the final segmentation prediction map obtained in step 1) Use the softmax method to normalize the prediction results Then the class-specific weights are calculated

[0027]

[0028] in, C represents the number of categories, represents the cth category of branch i; σ() represents the softmax function;

[0029] The final uncertainty weight is obtained by taking the average of the two network branches:

[0030]

[0031] Step 2.4) The final feature-level relational consistency loss is:

[0032]

[0033] Among them, W b1,c1 W c A subset of W c b1 sample.

[0034] The step 3) is specifically as follows: based on the predicted segmentation map of the two branches obtained in step 1) The pixel-level consistency loss of branches 1 and 2 is calculated by KL divergence:

[0035]

[0036]

[0037] in, Indicates that the second network branch constrains the first network branch. It means that the first network branch is used to constrain the second network branch, and τ represents the temperature factor hyperparameter.

[0038] The step 4) is specifically: based on the segmentation prediction map obtained in step 1) And its corresponding true label A, get the target task loss of the network branch

[0039]

[0040] Where i represents the i-th network branch, CE(.) and Dice(.) represent the cross entropy and Dice loss functions.

[0041] The step 5) is specifically: based on the feature level consistency loss L of step 2) fc , pixel-level consistency loss in step 3) And the task loss in step 4) Calculate the overall target loss L of branch i i :

[0042]

[0043] Beneficial effects of the present invention:

[0044] The Boundary Uncertainty based Relation Knowledgeable Distillation (BURKD) of the present invention mainly includes three steps: category-specific prototype generation, relational context modeling between samples and categories, and uncertainty-based reweighting. We integrated this method into a two-branch semantic segmentation network based on mutual learning, and explored the role of this method in the incomplete multimodal brain tumor segmentation scenario. In the present invention, we verified the effect of this method in the scenario of incomplete multimodal brain tumor segmentation model, surpassing the current SOTA and obtaining the best segmentation effect. The relational knowledge distillation method based on classification boundary uncertainty designed by the present invention fully takes into account the specificity of different instances in semantic segmentation. This enables the model to achieve SOTA results on the multimodal brain tumor semantic segmentation dataset (BraTs2018), surpassing other existing methods. BRIEF DESCRIPTION OF THE DRAWINGS

[0045] Figure 1 is a flow chart of the method of the present invention;

[0046] Figure 2 Schematic diagram of the BURKD algorithm flow of the present invention. DETAILED DESCRIPTION

[0047] The present invention is further described in detail below with reference to the accompanying drawings and specific embodiments.

[0048] like Figure 1 As shown, a relational knowledge distillation method based on category boundary uncertainty of the present invention comprises the following steps:

[0049] Step 1): An image passes through the feature encoders of two network branches at the same time Get the intermediate feature maps F1, F2;

[0050] Step 2): Then based on their respective decoders Get the final segmentation prediction map

[0051] Step 3): Based on the intermediate feature maps F1 and F2 obtained in step 1), calculate the feature-level consistency distillation loss L fc ;

[0052] The step 3) is specifically as follows: based on the intermediate feature maps F1 and F2 of the two branches obtained in step 1), the category-specific prototypes, the contextual relationships between modeled samples and categories, and the re-weighting of the uncertainty of category boundaries are calculated in sequence to obtain the feature-level consistency loss L fc , the specific operations are as in step 3.1) to step 3.3);

[0053] 3.1) Category-specific prototype generation: Given an intermediate feature map And its corresponding true label A∈R H×W×Z , calculate the class-specific prototypes of branches 1 and 2 respectively:

[0054]

[0055] in H i ,W i ,Z i Indicates the length, width, and height of the feature, d f Indicates the number of channels of the feature, H, W, Z represent the length, width, and height of the real label A, N c represents the number of pixels belonging to category c, Is an indicator function, which means that the value is 1 when the internal condition is True, otherwise it is 0;

[0056] 3.2) Relationship context modeling between samples and categories: Based on the category-specific prototypes obtained in step 3.1), the relationship maps R1 and R2 between samples and categories are obtained by the following formula:

[0057]

[0058] in, Represent the prototype of the c1th category of the b1th sample of branches 1 and 2, R1, R2∈R (B*C)×(B*C) , B represents the number of samples in a batch, and C represents the number of semantic categories.<a,b> Indicates the dot product operation of computing elements a and b, R i (b1, c1; b2, c2) represents the feature similarity between the c1 category of the b1th sample of the network branch i and the c2 category of the b2th sample. Then calculate the difference in the relationship distance between samples and categories between the two branches:

[0059]

[0060] in Represents the overall relationship difference between the b1th sample and c1th category between network branches 1 and 2;

[0061] 3.3) Re-weighting based on boundary uncertainty: In order to prioritize consistency learning on more challenging categories, we introduce a re-weighting method based on category boundary specificity. Based on the final segmentation prediction map obtained in step 2) Use the softmax method to normalize the prediction results Then the class-specific weights are calculated

[0062]

[0063] in C represents the number of categories, represents the cth category of branch i. The final uncertainty weight is the average of the two network branches:

[0064]

[0065] The final feature-level relation consistency loss;

[0066]

[0067] Step 4): Based on the segmentation prediction map obtained in step 2) Compute pixel-level consistency distillation loss The details are as follows:

[0068] Based on the predicted segmentation map of the two branches obtained in step 2) The pixel-level consistency loss of branches 1 and 2 is calculated by KL divergence:

[0069]

[0070] in Indicates that network branch 2 constrains branch 1. Indicates that network branch 1 constrains branch 2, and τ represents the temperature factor hyperparameter;

[0071] Step 5): Based on the dual-branch segmentation prediction map obtained in step 2) and its corresponding true category label A, and calculate the target loss of branches 1 and 2 respectively The details are as follows:

[0072] Based on the segmentation prediction map obtained in step 2) And its corresponding true label A, get the target task loss of the network branch

[0073]

[0074] Where i represents network branches 1 and 2, CE(.) and Dice(.) represent cross entropy and Dice loss functions;

[0075] Step 6): Combine the L obtained in step 3) fc and step 4) And the result obtained in step 5) Calculate the total optimization loss L i , as follows:

[0076] Based on the feature-level consistency loss L in step 3) fC , pixel-level consistency loss in step 4) And the task loss in step 5) Calculate the overall target loss L of branch i i :

[0077]

[0078] Step 7): Train the network model based on the BURKD algorithm, select the comprehensive loss as the final loss function, and use error back propagation to adjust the weights and biases of the network model. Select parameters based on the size of the Dice Similarity Coefficient (DSC), and finally determine a set of parameters that maximize the accuracy of brain tumor segmentation (DSC score).

[0079] The BURKD algorithm of the present invention can assist the knowledge distillation framework in knowledge transfer learning. The present invention applies it to brain tumor segmentation based on multimodal loss scenarios to verify the effectiveness of the algorithm. The model ultimately determines whether each pixel of the image belongs to its corresponding tumor category. According to the evaluation criteria, this method has significant effects on multi-category semantic segmentation.

[0080] The effects of the BURKD algorithm and some SOTA algorithms on the current BraTs2018 brain tumor segmentation dataset in the absence of any modality are as follows:

[0081]

[0082] Note:

[0083] (1) MCTSeg, M3AE, U-Net-MFI, RFNet, mmFormer and SMU-Net are the models that previously achieved SOTA on the BraTs2018 dataset, and BURKD is the model of the present invention;

[0084] (2) WT, TC, and ET represent Whole Tumor, Tumor Core, and Enhancing Tumor, respectively;

[0085] (3) The value in the table is the Dice Similarity Coefficient (DSC). The larger the value, the better the segmentation effect.

[0086] (4) Avg: represents the average DSC score of the three tumor categories. The larger the value, the better the segmentation effect;

[0087] (5) The bold black values ​​in the table indicate the best segmentation performance;

[0088] (6) Compared with the SOTA models MCTSeg, M3AE, U-Net-MFI, RFNet, mmFormer and SMU-Net, the three category segmentation indicators DSC on the same dataset of the present invention are improved by 0.5%, 1.1%, and 1.6%, respectively, with an average improvement of 1.1%.

Claims

1. A relational knowledge distillation method based on category boundary uncertainty, comprising the following steps: 1) Input an image into two network branches at the same time, first pass through the feature encoder Get the intermediate feature maps F1 and F2; then pass through the decoder Get the final segmentation prediction map 2) Based on the intermediate feature maps F1 and F2 obtained in step 1), calculate the feature-level consistency distillation loss L fc ; 3) Based on the segmentation prediction map obtained in step 1) Compute pixel-level consistency distillation loss 4) Based on the dual-branch segmentation prediction map obtained in step 1) and its corresponding true category label A, and calculate the target loss of the first network branch and the second network branch respectively 5) Combine the L obtained in step 2) fc and step 3) And the result obtained in step 5) Calculate the total optimization loss L i ; 6) Train the model based on the loss function of step 5); 7) Use the model trained in step 6) to obtain the segmentation or classification results of the image to be processed.

2. The relational knowledge distillation method based on category boundary uncertainty according to claim 1, characterized in that: In the step 1): The first network branch includes a feature encoder Decoder The second network branch includes a feature encoder Decoder 3. The relational knowledge distillation method based on category boundary uncertainty according to claim 1, characterized in that: The step 2) is specifically as follows: based on the intermediate feature maps F1 and F2 of the two branches obtained in step 1), the category-specific prototypes, the contextual relationships between modeled samples and categories, and the re-weighting of the uncertainty of category boundaries are calculated in sequence to obtain the feature-level consistency loss L fc , the specific operations are as follows: Step 2.1) Category-specific prototype generation: Given an intermediate feature map And its corresponding true label A∈R H×W×Z , respectively, calculate the class-specific prototypes of the first network branch and the second network branch: in, H i ,W i ,Z i Indicates the length, width, and height of the feature, d f Indicates the number of channels of the feature, H, W, Z represent the length, width, and height of the real label A, N c represents the number of pixels belonging to category c, Is an indicator function, which means that the value is 1 when the internal condition is True, otherwise it is 0; Step 2.2) Relationship context modeling between samples and categories: Based on the category-specific prototypes obtained in step 2.1), the relationship maps R1 and R2 between samples and categories are obtained by the following formula: in, Respectively represent the prototype of the c1th category of the b1th sample in the first network branch and the second network branch, R1, R2∈R (B*C)×(B*C) , B represents the number of samples in a batch, and C represents the number of semantic categories;<a,b> Indicates the dot product operation of computing elements a and b, R i (b1, c1; b2, c2) represents the feature similarity between the c1 category of the b1th sample and the c2 category of the b2th sample in the i-th network branch; then the difference in the relationship distance between samples and categories between the two branches is calculated: in, represents the overall relationship difference of the b1th sample and c1th category between the first network branch and the second network branch; Step 2.3) Reweighting based on boundary uncertainty: Based on the final segmentation prediction map obtained in step 1) Use the softmax method to normalize the prediction results Then the class-specific weights are calculated in, C represents the number of categories, represents the cth category of branch i; σ() represents the softmax function; The final uncertainty weight is obtained by taking the average of the two network branches: Step 2.4) The final feature-level relational consistency loss is: Among them, W b1,c1 W c A subset of W c b1 sample.

4. The relational knowledge distillation method based on category boundary uncertainty according to claim 1, characterized in that: The step 3) is specifically as follows: based on the predicted segmentation map of the two branches obtained in step 1) The pixel-level consistency loss of branches 1 and 2 is calculated by KL divergence: in, Indicates that the second network branch constrains the first network branch. It means that the first network branch is used to constrain the second network branch, and τ represents the temperature factor hyperparameter.

5. The relational knowledge distillation method based on category boundary uncertainty according to claim 1, characterized in that: The step 4) is specifically: based on the segmentation prediction map obtained in step 1) And its corresponding true label A, get the target task loss of the network branch Where i represents the i-th network branch, CE(.) and Dice(.) represent the cross entropy and Dice loss functions.

6. The relational knowledge distillation method based on category boundary uncertainty according to claim 1, characterized in that: The step 5) is specifically: based on the feature level consistency loss L of step 2) fc , pixel-level consistency loss in step 3) And the task loss in step 4) Calculate the overall target loss L of branch i i :

Images