A semi-supervised medical image segmentation method based on knowledge transfer joint training

Abstract

The application discloses a kind of semi-supervised medical image segmentation methods based on knowledge transfer joint training, comprising the following steps: step one, constructing training dataset D;Step two, set random sub-block fusion strategy;Step three, construct joint training framework;Step four, pre-training is carried out;Step five, self-training is carried out to obtain the final trained medical image segmentation model;Step six, using the final trained medical image segmentation model carries out medical image segmentation.The present application can implicitly increase sample representation space while reducing the distribution difference between labeled and unlabeled data representation by random sub-block fusion strategy, and can learn different knowledge through joint training framework, wherein the knowledge learned by one network can fill the cognitive deficiency of another network, and the two learn from each other, fill the cognitive upper limit of each other. And in the test phase, only one model is needed to participate in the test, without increasing any additional computational burden and cost.

Description

Technical Field

[0001] This invention mainly relates to the field of machine learning, specifically to a semi-supervised medical image segmentation method based on knowledge transfer joint training. Background Technology

[0002] In recent years, deep learning-based medical image segmentation methods have garnered widespread attention in the industry due to their outstanding performance. However, these fully supervised methods require extensive manual annotation, which is time-consuming, labor-intensive, and requires specialized domain knowledge; furthermore, manual annotation is subject to subjective bias. The high annotation cost limits the further development of deep learning methods in medical image segmentation. Therefore, to reduce annotation costs, semi-supervised segmentation methods have emerged. Semi-supervised segmentation methods only require a small portion of labeled data combined with a large amount of unlabeled data to train the segmentation model. However, current semi-supervised segmentation methods fail to utilize knowledge learned from labeled data to guide learning on unlabeled data. Consequently, much of the knowledge learned from reliably labeled data is wasted, causing the representations learned from unlabeled data to fail to align with those learned from labeled data in terms of probability distribution. Therefore, the current separate training approach of learning from labeled and unlabeled data limits the performance improvement of semi-supervised segmentation models. Summary of the Invention

[0003] To address the shortcomings of existing technologies, this invention discloses a semi-supervised medical image segmentation method based on knowledge transfer joint training.

[0004] The specific technical solution of the present invention is as follows:

[0005] A semi-supervised medical image segmentation method based on knowledge transfer joint training includes the following steps:

[0006] Step 1: Construct the training dataset D, D = D L ∪D U ;D L For a labeled dataset, D U This is an unlabeled dataset;

[0007] Step 2: Set the random sub-block merging strategy;

[0008] Step 3: Construct a joint training framework, which includes sub-network A and sub-network B; sub-network A is a V-Net network, and sub-network B is obtained by embedding attention layers in the second and third layers of the decoder of the V-Net network;

[0009] Step 4, Pre-training: Two sets of labeled data are randomly sampled from the labeled dataset. These two sets of labeled data are then fused using a random sub-block fusion strategy to obtain two mixed input data sets, which are then input into sub-networks A and B of the joint training framework for training. Simultaneously, the knowledge learned by sub-network B is unidirectionally transferred to sub-network A through KL divergence, resulting in pre-trained sub-networks A and B. The labels corresponding to the mixed input data are obtained by randomly fusion of the labels of the original labeled data to obtain two fused label data sets. The two fused label data sets supervise the corresponding outputs of sub-networks A and B, respectively.

[0010] Step 5: Self-training: Perform self-training on the pre-trained sub-network A and pre-trained sub-network B, and use the self-trained sub-network A as the final trained medical image segmentation model.

[0011] Step 6: Perform medical image segmentation using the final trained medical image segmentation model.

[0012] In a further improvement, the random sub-block fusion strategy in step two is as follows:

[0013] Generate a mask M∈{1} with the same size as the input data. H×W×D In M, the value at any spatial location is 1, and H, W, and D represent the height, width, and number of slices corresponding to the input data, respectively. Then, the input data is divided into N equal blocks along the H, W, and D dimensions, resulting in N... 3 A set of sub-masks i is the spatial index of the submask, and the submask... The value at any position in the array is also 1, and then from N... 3 Randomly select E submasks from the E submasks; transform the values ​​of the E submasks to 0. That is, the mixing mask is obtained. Submask representing a value of 0; mixed mask Reconstruction yields the hybrid mask M mixed ={0,1} H×W×D ;

[0014] By using a mixing mask, different input data X can be achieved. a X b Integration between them:

[0015] X ab =M mixed ⊙X a +(1-M mixed )⊙X b

[0016] Xba =M mixed ⊙X b +(1-M mixed )⊙X a

[0017] In the above formula, ⊙ represents the dot product, and X a and X b For different input data, X ab M is used for mixing masks mixed Zhongrongxi a Partial data, while utilizing (1-M) mixed Fusion X b The fused input data is obtained from a portion of the data; X ba For the mixing mask M mixed Zhongrongxi b Partial data, while utilizing (1-M) mixed Fusion X a The fused input data is obtained from a portion of the data.

[0018] A further improvement is made to the method for obtaining the two mixed input data in step four, as follows:

[0019] For labeled data and After fusion using a random sub-block fusion strategy, the mixed input data for the joint training framework is obtained. and The corresponding labeled data and The tags are also fused using a random sub-block fusion strategy to obtain fused tag data. and

[0020]

[0021]

[0022] in, and They are labeled data. and The tag.

[0023] In a further improvement, step four, where the pre-trained sub-network A and pre-trained sub-network B are obtained, is as follows: Let F... A (θ A Let θ be a subnetwork A with parameters θ. A F B (θ B Let θ be a subnetwork B with parameters θ. B Then the input data will be mixed. and The inputs are fed into subnetwork A and subnetwork B respectively, and the outputs are as follows:

[0024]

[0025]

[0026] Wherein, the output of subnetwork A Output of subnetwork B for The input is fed into the output corresponding to subnetwork A; for The input is sent to the output corresponding to subnetwork B;

[0027] The formula for unidirectionally transferring the knowledge learned by subnetwork B to subnetwork A is as follows:

[0028]

[0029] L KL Let KL represent the KL divergence loss function, and KL() represent the KL divergence function.

[0030] During the pre-training phase, the target loss function Loss is as follows:

[0031]

[0032] CE is the cross-entropy loss function. Dice is the Dice loss function; for For mixed input data The corresponding fusion tags, For mixed input data Corresponding fusion tags;

[0033] Pre-training for a preset number of training iterations yields pre-trained subnetworks A and B.

[0034] Further improvements, the specific steps of step five are as follows: [The text abruptly ends here, likely due to an incomplete sentence or a formatting error.] The inputs are fed into pre-trained sub-network A and pre-trained sub-network B respectively, to obtain the pseudo-labels output by pre-trained sub-network A and pre-trained sub-network B respectively. and

[0035] Will With labeled data tags The supervision signal of the pre-trained subnetwork A is obtained by fusing through a random sub-block fusion strategy. in

[0036]

[0037]

[0038] Will With labeled data tags The supervision information of the pre-trained subnetwork B is obtained by fusing through a random sub-block fusion strategy.

[0039]

[0040]

[0041] Then, for the labeled data and unlabeled data The input training data is obtained by fusion of random sub-blocks. and

[0042]

[0043]

[0044] The output of the pre-trained subnetwork A is represented as follows: The output of the pre-trained subnetwork B is represented as follows:

[0045] During self-training, the objective function L of the pre-trained subnetwork A is... A It is expressed as follows:

[0046]

[0047] in, Let the supervised objective function be denoted as:

[0048]

[0049] Unsupervised partial objective function Represented as

[0050]

[0051] Furthermore, λ represents unsupervised hyperparameter control. The proportion of the entire objective function, with a default value of 0.5;

[0052] The objective function L of the pre-trained subnetwork B B as follows

[0053]

[0054] Represented as:

[0055]

[0056] for Represented as

[0057]

[0058] During the self-training phase, the pre-trained sub-network A is supervised by the information of the pre-trained sub-network B, and the pre-trained sub-network B is supervised by the information of the pre-trained sub-network A. After self-training for a preset number of iterations, the self-trained sub-network A is used as the final trained medical image segmentation model.

[0059] The advantages and beneficial effects of this invention are as follows:

[0060] 1. Data knowledge level: A random sub-block fusion strategy is proposed to fuse unlabeled and labeled data. This strategy can implicitly increase the sample representation space while reducing the distribution difference between labeled and unlabeled data representations.

[0061] 2. Model Knowledge Level: A joint training framework is proposed, consisting of two slightly different sub-networks. One sub-network is the classic V-Net network, while the other sub-network, based on V-Net, embeds two attention layers in its decoder section. Apart from this, the two networks are identical. The purpose of this approach is to allow the two slightly different networks to learn different knowledge in a targeted manner. The knowledge learned by one network can fill the cognitive gaps of the other network, and vice versa. They learn from each other, filling each other's cognitive limitations. Attached Figure Description

[0062] Figure 1 A schematic diagram of the random sub-block fusion strategy

[0063] Figure 2 This diagram illustrates the pre-training phase. Dotted lines represent random sub-block fusion, and solid lines represent data flow. FA stands for V-Net, while FB embeds attention layers in the second and third layers of the decoder section, building upon V-Net. (Note: V-Net consists of an encoder and a decoder. Both the encoder and decoder contain five convolutional layers.)

[0064] Figure 3This diagram illustrates the self-training phase. Dotted lines represent random sub-block fusion, dashed lines represent data flows without gradient information, and solid lines represent data flows containing gradient information.

[0065] Figure 4 Visualization results of 10% annotation of the left atrium using different methods.

[0066] Figure 5 Visualization results of different methods based on 10% of the ACDC dataset annotations.

[0067] Figure 6 The image shows kernel density estimation curves for different methods. Each row represents a different method, and each column represents a comparison of kernel density estimation curves for different convolutional layer features. Detailed Implementation

[0068] The specific embodiments of the present invention will be described in detail below with reference to the accompanying drawings. It should be understood that the specific embodiments described herein are for illustration and explanation only and are not intended to limit the present invention.

[0069] First, let's introduce the basic concepts of semi-supervised segmentation methods. The input 3D medical image is X∈R. H×W×D H, W, and D represent the height, width, and number of slices in the input data, respectively. Our goal is to train a segmentation model F(θ), where θ is the model parameter. This model can classify any voxel in X, and its prediction result is defined as... C represents the total number of categories, where C = 0 represents the background (no segmentation needed), and C ≠ 0 represents the target category that needs segmentation. The entire dataset is defined as D = D0. L ∪D U Among them are labeled datasets N L This represents the number of labeled data and the number of unlabeled data. N U Let N be the number of unlabeled data points. L <<N U .

[0070] I. Random Sub-block Fusion Strategy.

[0071] First, generate a mask M∈{1} with the same size as the input data. H×W×D The value of M at any spatial location is 1. Then we divide it into N equal blocks along the H, W, and D dimensions to obtain N 3 Individual mask i is the spatial index of the submask, and the submask... The value at any position in N is also 1. Then from N... 3 E(E) is randomly selected from the submask. <N 3 ) number of masks Transform the values ​​of the E submasks to 0, that is... Based on the original spatial position index i of the submask, we obtain the hybrid mask. Finally, the hybrid mask M is reconstructed. mixed ={0,1} H×W×D By using a mixing mask, we have achieved different input data X a X b Integration between them:

[0072] X ab =M mixed ⊙X a +(1-M mixed )⊙X b (1)

[0073] X ba =M mixed ⊙X b +(1-M mixed )⊙X a (2)

[0074] In the above formula, ⊙ represents the dot product.

[0075] II. Joint Training Framework

[0076] This section is divided into a pre-training phase and a self-training phase. We will first introduce the pre-training phase.

[0077] Pre-training phase:

[0078] At this stage, model training is performed using only the labeled data portion. For example... Figure 2 As shown, according to formulas (1) and (2), the labeled data will be... (The superscript 'l' indicates that there is a label) The input training data is obtained by mixing according to the random sub-block fusion strategy. and The corresponding labels for the input training data can also be obtained through a random sub-block fusion strategy:

[0079]

[0080]

[0081] In the above formula, and It has labels with labeled data. F A (θ A Let θ be a subnetwork A with parameters θ. A F B (θ BLet θ be a subnetwork B with parameters θ. B Parameters are not shared between the two networks. Input data... and The inputs are fed into subnetwork A and subnetwork B respectively, and the outputs are obtained as follows:

[0082]

[0083]

[0084] The output of subnetwork A is represented as The superscript A indicates network F A The corresponding output. Similarly, the output of subnetwork B is represented as The superscript B indicates network F B The corresponding output. Considering that subnetwork B has an attention layer for enhancement, for the purpose of stable training, KL divergence is used to unidirectionally transfer the knowledge learned by subnetwork B to subnetwork A (Note: KL divergence is generally used to measure the "distance" between two probability distribution functions. Here, KL divergence is used to ensure that the outputs of the two networks have a certain diversity while the differences between them are not too large):

[0085]

[0086] Therefore, during the pre-training phase, the objective function is as follows:

[0087]

[0088] In the above formula, CE is the cross-entropy loss function. Dice is the Dice loss function.

[0089] Self-training phase:

[0090] like Figure 3 As shown, first, we will use unlabeled data The inputs are fed into pre-trained subnetworks A and B (dashed lines) respectively, resulting in pseudo-labels output by the two networks. and After random sub-block fusion (dotted and line parts), the supervisory signal of sub-network A Represented as

[0091]

[0092]

[0093] Similarly, for the supervision information of subnetwork B, the values ​​in equations (9) and (10) above will be used. Replace with That's all, you can get Then, for the labeled data and unlabeled data The input training data is obtained by fusion of random sub-blocks. and

[0094]

[0095]

[0096] Then the output of network A can be expressed as: And the output of network B is represented as The objective function is expressed as follows:

[0097]

[0098] In the above formula, Let the supervised objective function be denoted as:

[0099]

[0100] Unsupervised partial objective function Represented as

[0101]

[0102] Furthermore, λ represents unsupervised hyperparameter control. The proportion of the entire objective function, with a default value of 0.5.

[0103] For subnetwork B, we also have

[0104]

[0105] for It can be represented as

[0106]

[0107] for Represented as

[0108]

[0109] Subnetwork A is supervised by subnetwork B, while subnetwork B receives information supervision from subnetwork A. Given that subnetwork A is enhanced by subnetwork B during the pre-training phase, subnetwork A then feeds back into subnetwork B during the self-training phase. Furthermore, their different structures allow subnetwork B to continuously supervise subnetwork A from different perspectives, thereby achieving knowledge transfer at the network level.

[0110] The test results of this invention are as follows:

[0111] Segmentation results of a 3D MRI left atrium (LA) dataset. This dataset contains 100 MRI data points, of which 80 were used for training and 20 were used for testing and comparison. The results are as follows:

[0112] In our left atrial MRI segmentation experiment, we compared seven state-of-the-art semi-supervised methods, conducting two comparative experiments at 5% and 10% annotation levels.

[0113]

[0114] Note: UA-MT, SASSNet, DTC, URPC, MC-Net, SSNet, and BCPNet are all existing methods. SimKT is the method proposed in this invention, short for Simple Knowledge Transfer. Bold black text indicates the optimal result, and underlined text indicates the suboptimal result.

[0115] The above embodiments are only used to illustrate the technical solutions of the present invention, and are not intended to limit it. Although the present invention has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that modifications can still be made to the technical solutions described in the foregoing embodiments, or equivalent substitutions can be made to some or all of the technical features therein. Such modifications or substitutions do not cause the essence of the corresponding technical solutions to deviate from the scope of the technical solutions of the embodiments of the present invention.

Claims

1. A semi-supervised medical image segmentation method based on knowledge transfer joint training, characterized in that, Includes the following steps: Step 1: Construct the training dataset D, D = D L ∪D U ;D L For a labeled dataset, D U This is an unlabeled dataset; Step 2: Set the random sub-block merging strategy; Step 3: Construct a joint training framework, which includes sub-network A and sub-network B; sub-network A is a V-Net network, and sub-network B is obtained by embedding attention layers in the second and third layers of the decoder of the V-Net network; Step 4, Pre-training: Two sets of labeled data are randomly sampled from the labeled dataset. These two sets of labeled data are then fused using a random sub-block fusion strategy to obtain two mixed input data sets, which are then input into sub-networks A and B of the joint training framework for training. Simultaneously, the knowledge learned by sub-network B is unidirectionally transferred to sub-network A through KL divergence, resulting in pre-trained sub-networks A and B. The labels corresponding to the mixed input data are obtained by randomly fusion of the labels of the original labeled data to obtain two fused label data sets. The two fused label data sets supervise the corresponding outputs of sub-networks A and B, respectively. Step 5: Self-training: Perform self-training on the pre-trained sub-network A and pre-trained sub-network B, and use the self-trained sub-network A as the final trained medical image segmentation model. Step 6: Perform medical image segmentation using the final trained medical image segmentation model.

2. The semi-supervised medical image segmentation method based on knowledge transfer joint training as described in claim 1, characterized in that, In step two, the random sub-block fusion strategy is as follows: Generate a mask M∈{1} with the same size as the input data. H×W×D In M, the value at any spatial location is 1, and H, W, and D represent the height, width, and number of slices corresponding to the input data, respectively. Then, the input data is divided into N equal blocks along the H, W, and D dimensions, resulting in N... 3 A set of sub-masks i is the spatial index of the submask, and the submask... The value at any position in the array is also 1, and then from N... 3 Randomly select E submasks from the E submasks; transform the values ​​of the E submasks to 0. e = 0, 1, 2, ..., E, thus obtaining the mixing mask. Submask representing a value of 0; mixed mask Reconstruction yields the hybrid mask M mixed ={0,1} H×W×D ; By using a mixing mask, different input data X can be achieved. a X b Integration between them: X ab =M mixed ⊙X a +(1-M mixed )⊙X b X ba =M mixed ⊙X b +(1-M mixed )⊙X a In the above formula, ⊙ represents the dot product, and X a and X b For different input data, X ab M is used for mixing masks mixed Zhongrongx a Partial data, while utilizing (1-M) mixed Fusion X b The fused input data is obtained from a portion of the data; X ba For the mixing mask M mixed Zhongrongxi b Partial data, while utilizing (1-M) mixed ) Fusion x a The fused input data is obtained from a portion of the data.

3. The semi-supervised medical image segmentation method based on knowledge transfer joint training as described in claim 1, characterized in that, In step four, the method for obtaining the two mixed input data is as follows: For labeled data and After fusion using a random sub-block fusion strategy, the mixed input data for the joint training framework is obtained. and The corresponding labeled data and The tags are also fused using a random sub-block fusion strategy to obtain fused tag data. and in, and These are labeled data. and The tag.

4. The semi-supervised medical image segmentation method based on knowledge transfer joint training as described in claim 3, characterized in that, In step four, the steps to obtain the pre-trained sub-network A and pre-trained sub-network B are as follows: Let F A (θ A Let θ be a subnetwork A with parameters θ. A F B (θ B Let θ be a subnetwork B with parameters θ. B Then the input data will be mixed. and The inputs are fed into subnetwork A and subnetwork B respectively, and the outputs are as follows: Wherein, the output of subnetwork A Output of subnetwork B for The input is fed into the output corresponding to subnetwork A; for The input is sent to the output corresponding to subnetwork B; The formula for unidirectionally transferring the knowledge learned by subnetwork B to subnetwork A is as follows: L KL Let KL represent the KL divergence loss function, and KL() represent the KL divergence function. During the pre-training phase, the target loss function Loss is as follows: CE is the cross-entropy loss function. Dice is the Dice loss function; for For mixed input data The corresponding fusion tags, For mixed input data Corresponding fusion tags; Pre-training for a preset number of training iterations yields pre-trained subnetworks A and B.

5. The semi-supervised medical image segmentation method based on knowledge transfer joint training as described in claim 1, characterized in that, The specific steps of step five are as follows: Transfer unlabeled data... The inputs are fed into pre-trained sub-network A and pre-trained sub-network B respectively, to obtain the pseudo-labels output by pre-trained sub-network A and pre-trained sub-network B respectively. and Will With labeled data tags The supervision signal of the pre-trained subnetwork A is obtained by fusing through a random sub-block fusion strategy. in Will With labeled data tags The supervision information of the pre-trained subnetwork B is obtained by fusing through a random sub-block fusion strategy. Then, for the labeled data and unlabeled data The input training data is obtained by fusion of random sub-blocks. and The output of the pre-trained subnetwork A is represented as follows: The output of the pre-trained subnetwork B is represented as follows: During self-training, the objective function L of the pre-trained subnetwork A is... A It is expressed as follows: in, Let the supervised objective function be denoted as: Unsupervised partial objective function Represented as Furthermore, λ represents unsupervised hyperparameter control. The proportion of the entire objective function, with a default value of 0.5; The objective function L of the pre-trained subnetwork B B as follows Represented as: for Represented as During the self-training phase, the pre-trained sub-network A is supervised by the information of the pre-trained sub-network B, and the pre-trained sub-network B is supervised by the information of the pre-trained sub-network A. After self-training for a preset number of iterations, the self-trained sub-network A is used as the final trained medical image segmentation model.

Images