WO2024221231A1 - Semi-supervised learning method and apparatus based on model framework - Google Patents

Abstract

The present invention relates to a semi-supervised learning method and apparatus based on a model framework. The method and apparatus comprise: providing an overall network framework, wherein the overall network framework comprises a student model, a teacher model, a projector network, and an output layer network; inputting medical images into the student model and the teacher model, and obtaining projected feature representations from the outputs of the student model and the teacher model by means of the projector network; inputting the projected feature representations into the output layer network to obtain a final segmentation result of the student model and the teacher model; and designing a network loss function, and training the overall network framework by using the network loss function. In the present invention, a Mean-Teacher model is used, and contrastive learning as a loss function is added into a teacher-student model, so that the accuracy of consistency learning is further improved, the segmentation accuracy is thus improved, and a more accurate segmentation result is finally obtained.

Description

A semi-supervised learning method and device based on model framework Technical Field

The present invention relates to the field of medical image segmentation, and in particular to a semi-supervised learning method and device based on a model framework.

Background Art

The structure of the left atrium is important information for clinicians to diagnose and treat atrial fibrillation, the most common heart rhythm disorder. Medical image segmentation is the basis of various medical image applications, such as determining cancer staging, formulating treatment plans, radiomics analysis, and developing personalized medical services. Among the segmentation tasks, tumor target delineation is a key step in cancer treatment, the purpose of which is to maximize the concentration of radioactive agents in the target area and minimize or even avoid damage to surrounding normal tissues and organs. However, manual delineation of tumor targets is a time-consuming and laborious process, and the accuracy of manual annotation depends largely on the experience and knowledge of oncologists. Differences between different doctors may also lead to different annotations for the same tumor. Supervised 3D medical image segmentation methods have achieved great success, but they rely on a large amount of labeled data, which greatly limits the scope of application of supervised methods. Semi-supervised segmentation methods solve this problem by using a large amount of unlabeled data and a small amount of labeled data. At present, the most successful semi-supervised learning method is based on consistency learning, which minimizes the distance between model responses obtained from perturbed views of unlabeled data. In addition, contrastive learning has been proven to be an effective unsupervised learning method. Therefore, it is important to study and develop a semi-supervised segmentation method for medical images based on contrastive learning, while ensuring the segmentation accuracy and minimizing the dependence on labeled data, which has important scientific significance and application prospects in the field of medical diagnosis.

Lequan Yu et al. published an article titled "Uncertainty-aware Self-ensembling Model for Semi-supervised 3D Left Atrium Segmentation" at the MICCAI conference in 2019. This method proposes an uncertainty-based semi-supervised learning framework by combining the Mean-Teacher model and Monte Carlo simulation. The student model gradually learns from meaningful and reliable targets by leveraging the uncertainty information of the teacher model. In addition to generating target outputs, the teacher model also estimates the uncertainty of each target prediction through Monte Carlo Dropout. Guided by the estimated uncertainty, unreliable predictions are filtered out when calculating the consistency loss, and only reliable predictions (low uncertainty) are retained. As a result, the student model is optimized and receives more reliable supervision, which in turn encourages the teacher model to generate higher quality targets.

Ting Chen et al. published the article "SimCLR: A Simple Framework for Contrastive Learning of Visual Representations" at ICML in 2020. The article proposed a new self-supervised contrastive learning method. The SimCLR learning framework mainly consists of four components, including a random data augmentation module, a feature encoding module, a feature projection module, and a contrastive loss module. The core idea is to maximize the consistency between different augmented views of the same data example to learn representations.

In summary, the prior art has the following technical defects:

1. Medical image data annotation is difficult, time-consuming and laborious;

2. Manually labeled data depends on expert experience and knowledge, and different experts have differences;

3. Since most algorithms are designed based on a few parts, the algorithm robustness is poor.

Summary of the invention

The embodiment of the present invention provides a semi-supervised learning method and device based on a model framework to ultimately obtain a more accurate medical image segmentation result.

According to an embodiment of the present invention, a semi-supervised learning method based on a model framework is provided, comprising the following steps:

S101: Setting up the overall network framework, which includes a student model, a teacher model, a projector network, and an output layer network;

S102: inputting the medical image into the student model and the teacher model, and passing the output of the student model and the teacher model through the projector network to obtain a projection feature representation;

S103: Input the projected feature representation to the output layer network to obtain the final segmentation results of the student model and the teacher model;

S104: Design a network loss function, and use the network loss function to train the overall network framework.

Furthermore, both the student model and the teacher model use V-Net as the backbone network, and the network encoder The network and decoder contain 4 convolution-pooling layers respectively. The convolution kernel of the convolution layer is 3x3x3, the convolution kernel of the pooling layer is 2x2x2, the output channels are 16, 32, 64, 128, 128, 64, 32, and 16 respectively, and the activation function is ReLU.

Furthermore, the projector network contains two convolutional layers, the output channels of the first convolutional layer are 16, the output channels of the second convolutional layer are 8, and the convolution kernel size is 3x3x3.

Furthermore, the output layer network is a convolutional layer, the input is the output of V-Net, the number of channels is 16, the output channel of the output layer network is 2, and the convolution kernel size is 1x1x1.

Furthermore, the network loss function is divided into four parts, which is the sum of the student model supervised loss, the student-teacher model consistency loss, the student-teacher model cross loss, and the student-teacher model contrast loss.

Furthermore, for the supervised loss of the student model, given a training dataset D = {(x ₁ , y ₁ ), (x ₂ , y ₂ ), …, (x _N , y _N )}, where x = {x ₁ , x ₂ , …, x _N } is the network input image, y = {y ₁ , y ₂ , …, y _N } is the image annotated by the doctor, and N is the total number of training samples; the student supervised loss uses Dice loss and cross entropy loss, expressed as:

in Represents the prediction result of the student model for labeled data, and ε is a small constant.

Furthermore, the consistency loss and cross loss of the student model and the teacher model both use mean square error loss, expressed as:

y _1i and are the output results of the student model and the teacher model respectively.

Furthermore, the contrast loss between the student model and the teacher model is expressed as:

in is an indicative function, whose value is 1 if and only if k≠i, otherwise it is 0; τ is a constant; is a cosine similarity function; z _i and z _j are the projection output results of the student model and the teacher model respectively.

Furthermore, the Adam optimizer is used to optimize the loss function.

According to another embodiment of the present invention, a semi-supervised learning device based on a model framework is provided, comprising:

The framework setting unit is used to set the overall network framework, which includes the student model, the teacher model, the projector network and the output layer network;

A projection feature representation acquisition unit is used to input the medical image into the student model and the teacher model, and pass the output of the student model and the teacher model through the projector network to obtain the projection feature representation;

The final segmentation result acquisition unit is used to input the projection feature representation into the output layer network to obtain the final segmentation results of the student model and the teacher model;

The network loss function design unit is used to design the network loss function and use the network loss function to train the overall network framework.

A storage medium stores a program file capable of implementing any of the above-mentioned semi-supervised learning methods based on a model framework.

A processor is used to run a program, wherein when the program is run, any of the above-mentioned semi-supervised learning methods based on the model framework is executed.

The semi-supervised learning method and device based on the model framework in the embodiment of the present invention designs a semi-supervised learning framework by considering self-supervised learning and contrastive learning methods. In combination with the Mean-Teacher model, contrastive learning is added as a loss function to the teacher-student model to further increase the accuracy of consistency learning, thereby improving the segmentation accuracy performance. The present invention ultimately obtains a more accurate segmentation result.

BRIEF DESCRIPTION OF THE DRAWINGS

FIG1 is a framework diagram of a semi-supervised segmentation network based on contrastive learning of the present invention;

FIG. 2 is an experimental verification diagram of the present invention.

DETAILED DESCRIPTION

In order to enable those skilled in the art to better understand the scheme of the present invention, the technical scheme in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work should fall within the scope of protection of the present invention.

It should be noted that the terms "first", "second", etc. in the specification and claims of the present invention and the above-mentioned drawings are used to distinguish similar objects, and are not necessarily used to describe a specific order or sequence. It should be understood that the data used in this way can be interchanged where appropriate, so that the embodiments of the present invention described herein can be implemented in an order other than those illustrated or described herein. In addition, the terms "including" and "having" and any variations thereof are intended to cover non-exclusive inclusions, for example, a process, method, system, product or device that includes a series of steps or units is not necessarily limited to those steps or units that are clearly listed, but may include other steps or units that are not clearly listed or inherent to these processes, methods, products or devices.

In order to solve the problem of insufficient annotated data in medical image segmentation, the present invention considers self-supervised learning and contrastive learning methods, designs a semi-supervised learning framework, combines the Mean-Teacher model to add contrastive learning as a loss function to the teacher-student model, further increases the accuracy of consistency learning, and thus improves the segmentation accuracy performance. The present invention ultimately obtains a more accurate segmentation result. The present invention proposes a semi-supervised learning method and device based on a model framework, which combines consistency learning and contrastive learning, and uses mean square error loss to alternately optimize the model results, ultimately obtaining a more accurate segmentation result.

To achieve the above purpose, the specific operation steps of the technical solution adopted by the present invention are as follows:

Step 1: Set up the overall network framework

The network adopts the teacher-student model framework as shown in Figure 1. It mainly includes the student model, the teacher Teacher model, projector network, and output layer network.

Both the student model and the teacher model use V-Net as the backbone network. The network encoder and decoder contain 4 convolution-pooling layers respectively. The convolution kernel of the convolution layer is 3x3x3, and the convolution kernel of the pooling layer is 2x2x2. The output channels are 16, 32, 64, 128, 128, 64, 32, and 16 respectively, and the activation function uses ReLU.

Table 1 V-Net network parameter settings

Step 2: Set up the comparative projector network

First, the output of the student-teacher model is passed through the projector network to obtain the projected feature representation. The projector network contains two convolutional layers, the output channels of the first convolutional layer are 16, the output channels of the second convolutional layer are 8, and the convolution kernel size is 3x3x3.

Step 3: Set up the output layer network

As shown in Figure 1, the output layer network is a convolutional layer, the input is the output of V-Net, the number of channels is 16, the output channel of the output layer network is 2, the convolution kernel size is 1x1x1, and the output is the final segmentation result of the student-teacher model.

Step 4: Design network loss function

The entire network loss function is divided into four parts: the labeled supervised loss of the student model, the consistency loss of the student-teacher model, the cross loss of the student-teacher model, and the contrast loss of the student-teacher model.

For the supervised loss of the student model, given the training dataset D = {(x ₁ ,y ₁ ), (x ₂ ,y ₂ ), …, (x _N ,y _N )}, where x = {x ₁ ,x ₂ , …,x _N } is the network input image, y = {y ₁ ,y ₂ , …,y _N } is the image annotated by the doctor, and N is the total number of training samples. The student supervised loss uses Dice loss and cross entropy loss, expressed as:

in It represents the prediction result of the student model for the labeled data. ε is a very small constant to avoid the denominator being 0. It is set to 0.001 in the experiment.

The consistency loss and cross loss of the student-teacher model both use mean square error loss, expressed as:

y _1i and are the output results of the student model and the teacher model respectively.

The student-teacher model contrast loss is expressed as:

in is an indicative function whose value is 1 if and only if k≠i, and 0 otherwise. τ is a constant set to 2. is a cosine similarity function. _{z i} and z _j are the projection output results of the student model and the teacher model respectively.

The final loss function of the network is the sum of the above four parts.

Step 5: For the overall network framework designed above, use Adam optimizer to optimize the loss function.

Step 6: Train the network.

Compared with the prior art, the present invention has the following beneficial effects: the present invention increases the cross loss by considering the consistency between the student-teacher model, so that the results of the two models are mutually labeled, High semi-supervised learning accuracy. In addition, an unsupervised learning mechanism such as contrastive learning is introduced into the network, which further improves the accuracy of segmentation with the assistance of a portion of labeled data. The present invention effectively improves the performance of network-based supervised segmentation, and the segmentation results are better.

The present invention is verified using MRI data, and the present invention is proven to be feasible through experiments, simulations, and use, and the experimental results are shown in Figure 2. In addition to MRI data, the present invention can also be applied to other modal medical image data such as CT and PET.

Example 2

A storage medium stores a program file capable of implementing any of the above-mentioned semi-supervised learning methods based on a model framework.

Example 3

A processor is used to run a program, wherein when the program is run, any of the above-mentioned semi-supervised learning methods based on the model framework is executed.

The serial numbers of the above embodiments of the present invention are only for description and do not represent the advantages or disadvantages of the embodiments.

In the above embodiments of the present invention, the description of each embodiment has its own emphasis. For parts that are not described in detail in a certain embodiment, reference can be made to the relevant descriptions of other embodiments.

In the several embodiments provided in this application, it should be understood that the disclosed technical content can be implemented in other ways. Among them, the system embodiments described above are only schematic. For example, the division of units can be a logical function division. There may be other division methods in actual implementation. For example, multiple units or components can be combined or integrated into another system, or some features can be ignored or not executed. Another point is that the mutual coupling or direct coupling or communication connection shown or discussed can be through some interfaces, indirect coupling or communication connection of units or modules, which can be electrical or other forms.

The units described as separate components may or may not be physically separated, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed on multiple units. Some or all of the units may be selected according to actual needs to achieve the purpose of the present embodiment.

In addition, each functional unit in each embodiment of the present invention may be integrated into one processing unit, or each unit may exist physically separately, or two or more units may be integrated into one unit. The above-mentioned integrated unit may be implemented in the form of hardware or in the form of software functional units.

If the integrated unit is implemented in the form of a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, in essence, or the part that contributes to the prior art, or all or part of the technical solution can be embodied in the form of a software product. The computer software product is stored in a storage medium, including several instructions for a computer device (which can be a personal computer, server or network device, etc.) to perform all or part of the steps of the methods of each embodiment of the present invention. The aforementioned storage medium includes: U disk, read-only memory (ROM, Read-Only Memory), random access memory (RAM, Random Access Memory), mobile hard disk, magnetic disk or optical disk and other media that can store program codes.

The above is only a preferred embodiment of the present invention. It should be pointed out that for ordinary technicians in this technical field, several improvements and modifications can be made without departing from the principle of the present invention. These improvements and modifications should also be regarded as the scope of protection of the present invention.

Claims (10)

A semi-supervised learning method based on a model framework, characterized by comprising the following steps: S101: Setting up the overall network framework, which includes a student model, a teacher model, a projector network, and an output layer network; S102: inputting the medical image into the student model and the teacher model, and passing the output of the student model and the teacher model through the projector network to obtain a projection feature representation; S103: Input the projected feature representation to the output layer network to obtain the final segmentation results of the student model and the teacher model; S104: Design a network loss function, and use the network loss function to train the overall network framework. The semi-supervised learning method based on the model framework according to claim 1 is characterized in that both the student model and the teacher model use V-Net as the backbone network, the network encoder and decoder respectively include 4 convolution-pooling layers, the convolution kernel of the convolution layer is 3x3x3, the convolution kernel of the pooling layer is 2x2x2, the output channels are 16, 32, 64, 128, 128, 64, 32, and 16 respectively, and the activation function uses ReLU. The semi-supervised learning method based on the model framework according to claim 1 is characterized in that the projector network comprises two convolutional layers, the output channel of the first convolutional layer is 16, the output channel of the second convolutional layer is 8, and the convolution kernel size is 3x3x3. The semi-supervised learning method based on the model framework according to claim 1 is characterized in that the output layer network is a convolutional layer, the input is the output of V-Net, the number of channels is 16, the output channel of the output layer network is 2, and the convolution kernel size is 1x1x1. The semi-supervised learning method based on the model framework according to claim 1 is characterized in that the network loss function is divided into four parts, which is the sum of the student model supervised loss, the student-teacher model consistency loss, the student-teacher model cross loss and the student-teacher model contrast loss. The semi-supervised learning method based on the model framework according to claim 5 is characterized in that, for the supervised loss of the student model, given a training data set D = {(x ₁ , y ₁ ), (x ₂ , y ₂ ), …, (x _N , y _N )}, wherein x = {x ₁ , x ₂ , …, x _N } is the network input image, y = {y ₁ , y ₂ , …, y _N } is the image annotated by the doctor, and N is the total number of training samples; the student supervision loss uses Dice loss and cross entropy loss, expressed as:

in Represents the prediction result of the student model for labeled data, and ε is a small constant. The semi-supervised learning method based on the model framework according to claim 5 is characterized in that the consistency loss and the cross loss of the student model and the teacher model both use the mean square error loss, which is expressed as:
y _1i and are the output results of the student model and the teacher model respectively. The semi-supervised learning method based on the model framework according to claim 5 is characterized in that the contrast loss between the student model and the teacher model is expressed as:
in is an indicative function, whose value is 1 if and only if k≠i, otherwise it is 0; τ is a constant; is a cosine similarity function; z _i and z _j are the projection output results of the student model and the teacher model respectively. The semi-supervised learning method based on the model framework according to claim 1 is characterized in that an Adam optimizer is used to optimize the loss function. A semi-supervised learning device based on a model framework, characterized by comprising: The framework setting unit is used to set the overall network framework, which includes the student model, the teacher model, the projector network and the output layer network; A projection feature representation acquisition unit is used to input the medical image into the student model and the teacher model, and pass the output of the student model and the teacher model through the projector network to obtain the projection feature representation; The final segmentation result acquisition unit is used to input the projection feature representation into the output layer network to obtain the final segmentation results of the student model and the teacher model; The network loss function design unit is used to design the network loss function and use the network loss function to train the overall network framework.