Training method and system of medical image processing model, equipment and medium

Abstract

The invention discloses a medical image processing model training method, system, device and medium, and relates to the technical field of model training, and the method comprises the steps: obtaining a plurality of unlabeled medical images, and constructing a plurality of first image segmentation models and a second image segmentation model; inputting the unmarked medical image into a plurality of first image segmentation models for self-supervised pre-training; migrating the self-supervised training representation of the first image segmentation model to a second image segmentation model; and constructing an online path and an offline path in the migrated second image segmentation model, wherein the second image segmentation model in the online path is used for final downstream task testing. According to the adaptive method for self-supervised pre-training and downstream task fine tuning provided by the invention, the problem of disastrous forgetting of self-supervised representation can be relieved while the second image segmentation model learns the self-supervised representation of the first image segmentation model.

Description

Technical Field

[0001] The present invention relates to the field of model training technology, and in particular to a training method, system, equipment and medium for a medical image processing model. Background Art

[0002] In recent years, deep learning-based methods have achieved great success in 3D image analysis. In order to ensure that the trained model has high accuracy and generalization, a large amount of labeled data is generally required. However, in medical image analysis, the labeling of 3D images is very time-consuming and labor-intensive, and requires the professional knowledge of clinicians.

[0003] In comparison, unlabeled 3D images are relatively easier to obtain and the data is larger. In recent years, self-supervised learning has provided a solution to reduce the cost of labeling. It can learn visual representations without manual labeling. For example, the paper "Ibot: Image Bert Pre-training with Online Tokenizer" published in the conference "The 10th International Conference on Learning Representations" is a self-supervised learning method that makes the input mask image consistent with the information obtained by the model under the constraints of the teacher-student model, and the feature extractor trained by this method has shown excellent performance in downstream tasks such as image segmentation, classification and detection.

[0004] However, current self-supervised learning methods usually fine-tune the first image segmentation model pre-trained upstream directly on small downstream data, while for downstream tasks with a small amount of labeled data, the second image segmentation model is often more suitable. In addition, when fine-tuning the first image segmentation model on downstream tasks, the catastrophic forgetting problem of the upstream pre-trained model representation usually occurs, that is, the fine-tuned first image segmentation model gradually overfits on the downstream tasks, which will cause the accuracy and generalization of the fine-tuned model to be greatly reduced. Summary of the invention

[0005] The purpose of the present invention is to address the deficiencies of the above-mentioned prior art and provide a training method, system, device and medium for a medical image processing model to solve the problem of catastrophic forgetting of the representation of the upstream pre-trained model that usually occurs when fine-tuning the first image segmentation model on the downstream task in the prior art.

[0006] The present invention specifically provides the following technical solutions:

[0007] A training method for a medical image processing model comprises the following steps:

[0008] Acquire a number of unlabeled medical images and labeled medical images, and construct a number of first image segmentation models and a second image segmentation model; wherein the first image segmentation model has a higher number of parameters than the second image segmentation model;

[0009] Performing self-supervised pre-training on a plurality of first image segmentation models through the unlabeled medical images to obtain a self-supervised training representation with a minimized prediction distribution, and migrating the self-supervised training representation to a second image segmentation model;

[0010] Constructing an online path and an offline path, and inputting the labeled medical image into a second image segmentation model in the online path and the offline path to obtain a loss constraint of a downstream task and a prediction constraint of an offline path output;

[0011] The parameters of the second image segmentation model are adjusted through the loss constraint and the prediction constraint, and the trained medical image processing model is obtained by combining the second image segmentation model after parameter adjustment with the first image segmentation model after self-supervised pre-training.

[0012] Preferably, performing self-supervised pre-training on a plurality of first image segmentation models using the unlabeled medical images comprises the following steps:

[0013] The standard teacher-student paradigm is used as self-supervised pre-training; the goal of self-supervised pre-training is to minimize the prediction distribution of the teacher model and the student model;

[0014] During the self-supervised pre-training process, the parameters of the teacher model are updated by the exponential moving average of the student model parameters, and the parameters of the student model are updated by minimizing the output prediction of the teacher model to obtain a self-supervised training representation with minimized prediction distribution.

[0015] Preferably, migrating the self-supervised training representation to a second image segmentation model comprises the following steps:

[0016] In the first image segmentation model F (#,i) and the second image segmentation model F S Then add the same number of mappers to map the encoder output to the feature space;

[0017] Given an unlabeled medical image training sample x, use the data augmentation module on the medical image training sample x Obtain the augmented image blocks x1 and x2;

[0018] The x1 and x2 are respectively input into a pre-trained first image segmentation model and a randomly initialized second image segmentation model, and the features of the first image segmentation model and the second image segmentation model are aligned by introducing a consistency loss.

[0019] Preferably, the data augmentation module is used Obtain the augmented image blocks x1 and x2, the specific expression is:

[0020] x1=τ1(x),x2=τ2(x)

[0021] Among them, x1 and x2 are input into F (#,i) and F S middle.

[0022] Preferably, the step of aligning the features of the first image segmentation model and the second image segmentation model by introducing consistency loss comprises the following steps:

[0023] The loss function in the consistency loss is:

[0024]

[0025] Among them, L mimic (·) is the loss function, They represent the first image segmentation models with different parameters pre-trained in three upstreams, e is an exponential function, t1 and t2 are temperature coefficients, σ ​​is a softMax function, Yes F S Mapper added later, P # Yes F (#,i) Mapper added later.

[0026] Preferably, the online path and the offline path are constructed, and the specific expression is:

[0027] W off =α*W off +(1-α)*W on

[0028] Among them, W off and W on Respectively and The parameter weights, is the offline second image segmentation model, is the second online image segmentation model, The parameters are updated by The parameters are implemented as exponential moving average, and α is the momentum coefficient.

[0029] Preferably, adjusting the parameters of the second image segmentation model by using the loss constraint and the prediction constraint comprises the following steps:

[0030] For online paths By minimizing the downstream task loss and minimizing The output distribution difference is The parameter update L FT , the specific expression is:

[0031]

[0032] Among them, x l is the labeled data in the downstream task, L DS is the loss function of the downstream specific task, σ is the softMax function, and β is the loss coefficient used to adjust the impact of the two losses.

[0033] The present invention provides a training system for a medical image processing model, comprising:

[0034] An acquisition module, used to acquire a number of unlabeled medical images and a number of labeled medical images, and acquire a number of first image segmentation models and a second image segmentation model; wherein the complexity of the first image segmentation model is higher than that of the second image segmentation model;

[0035] A training module, configured to perform self-supervised pre-training on a plurality of first image segmentation models using the unlabeled medical images, obtain a self-supervised training representation with a minimized prediction distribution, and transfer the self-supervised training representation to a second image segmentation model;

[0036] A testing module, used for constructing an online path and an offline path, and inputting the labeled medical image into a second image segmentation model in the online path and the offline path to obtain a loss constraint of a downstream task and a prediction constraint of an offline path output;

[0037] The model generation module is used to adjust the parameters of the second image segmentation model through the loss constraint and the prediction constraint, and obtain the trained medical image processing model by combining the second image segmentation model after parameter adjustment with the first image segmentation model after self-supervised pre-training.

[0038] The present invention provides a computer device, including a memory and a processor, wherein a program is stored in the memory, and when the program is executed by the processor, the processor executes the steps of the above-mentioned training method of a medical image processing model.

[0039] The present invention provides a storage medium having a computer program stored thereon, characterized in that when the computer program is executed by a processor, the steps of the training method of the above-mentioned medical image processing model are implemented.

[0040] Compared with the prior art, the present invention has the following significant advantages:

[0041] The present invention transfers the self-supervised representation learned by the first image segmentation model to the second image segmentation model, so that the second image segmentation model has the representation ability of the first image segmentation model, and constructs an online path and an offline path when fine-tuning the second image segmentation model on downstream small data. The parameters of the second image segmentation model are adjusted through the prediction constraints of the discrete path output and the loss constraints of the downstream tasks, thereby alleviating the catastrophic forgetting problem of the pre-trained self-supervised representation and greatly enhancing the accuracy and generalization of the medical image processing model after fine-tuning training. BRIEF DESCRIPTION OF THE DRAWINGS

[0042] Figure 1 It is the overall structure diagram of the present invention. DETAILED DESCRIPTION

[0043] The following is a clear and complete description of the technical solutions of the embodiments of the present invention in conjunction with the drawings in the present invention. Obviously, the described embodiments are 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.

[0044] Reference Figure 1 The adaptation scheme of self-supervised pre-training and downstream task fine-tuning includes self-supervised pre-training of the upstream first image segmentation model, migration of the upstream first image segmentation model to the pre-trained representation of the second image segmentation model, and fine-tuning of the downstream second image segmentation model. Among them, the first image segmentation model is a large model and the second image segmentation model is a small model. First, the pre-training of the three first image segmentation models is completed on a large number of unlabeled samples using the teacher-student self-supervised paradigm, and then the parameters after pre-training are fixed. Given an unlabeled training sample x, use the data augmentation module on it The augmented image blocks are x1 and x2, which are input into the pre-trained first image segmentation model and the randomly initialized second image segmentation model respectively. The features of the first image segmentation model and the second image segmentation model are aligned by introducing consistency loss to achieve the migration of the pre-trained representation of the first image segmentation model to the second image segmentation model. For downstream small data tasks, the second image segmentation model in the online path fine-tunes the parameters through the loss constraints of the downstream tasks and the prediction constraints of the output of the second image segmentation model in the offline path. After fine-tuning, the second image segmentation model in the online path is directly used for the final downstream task test.

[0045] The present invention provides a training method for a medical image processing model, which specifically comprises the following steps:

[0046] Step S1: obtaining a number of unlabeled medical images and a number of labeled medical images, and constructing a number of first image segmentation models and a second image segmentation model, wherein the complexity and parameter quantity of the first image segmentation model are higher than those of the second image segmentation model.

[0047] Specifically:

[0048] Upstream first image segmentation model self-supervised pre-training:

[0049] The standard teacher-student paradigm is used as self-supervised pre-training, and the pre-training goal is to minimize the prediction distribution of the teacher model and the student model. The teacher model and the student model have the same network structure, and the model weights are randomly initialized.

[0050] During the self-supervised pre-training process, the parameters of the teacher model are updated by the exponential moving average of the student model parameters, and the parameters of the student model are updated by minimizing the output prediction of the teacher model to obtain a self-supervised training representation with minimized prediction distribution.

[0051] Step S2: self-supervised pre-training is performed on several first image segmentation models through unlabeled medical images to obtain a self-supervised training representation with minimized prediction distribution, and the self-supervised training representation is transferred to the second image segmentation model.

[0052] Specifically:

[0053] Migrate the pre-trained representation of the first image segmentation model to the second image segmentation model:

[0054] Considering that the second image segmentation model has better performance than the first image segmentation model in downstream small data tasks, after pre-training multiple first image segmentation models F using upstream unlabeled data (#,i) After that, it is necessary to transfer the self-supervised pre-trained representation of the first image segmentation model to the second image segmentation model F S middle.

[0055] First, in F (#,i) and F S The mapper is added later, and its basic structure is convolution layer-batch normalization layer-activation function layer-convolution layer, which is used to map the output of the encoder to the feature space. S With F (#,i) Same number of mappers For learning the corresponding F (#,i) Specifically, given an unlabeled training sample x, a data augmentation module is used on the medical image training sample x And from the data augmentation module Randomly select two different augmentation methods The two image blocks extracted from x are augmented, and the augmented image blocks are represented as x1 and x2.

[0056] x1=τ1(x),x2=τ2(x)

[0057] Among them, x1 and x2 are respectively input into the pre-trained first image segmentation model and the randomly initialized second image segmentation model. In order to migrate the pre-trained representation of the first image segmentation model, the consistency loss is introduced to align the features of the first image segmentation model and the second image segmentation model.

[0058] The loss function in consistency loss is as follows:

[0059]

[0060] Among them, L mimic (·) is the loss function, They represent the first image segmentation models with different parameters pre-trained in three upstream models. e is an exponential function, t1 and t2 are temperature coefficients, and σ is a softMax function. This loss function minimizes the feature mapping distance between the first image segmentation model and the second image segmentation model, so that the second image segmentation model can learn the rich representation ability of the first image segmentation model. Yes F S Mapper added later, P # Yes F (#,i) Mapper added later.

[0061] Step S3: construct an online path and an offline path, and input the labeled medical image into the second image segmentation model in the online path and the offline path to obtain the loss constraint of the downstream task and the prediction constraint of the offline path output.

[0062] Specifically:

[0063] Downstream second image segmentation model fine-tuning:

[0064] After migrating the self-supervised representation of the first image segmentation model to the second image segmentation model, in order to alleviate the catastrophic forgetting problem of the upstream representation, the present invention constructs an online path and an offline path in the downstream task fine-tuning. The two paths initially have the same model weights. and in, The parameters are updated by The parameters are implemented as exponential moving average. Construct online path and offline path, the specific expression is:

[0065] W off =α*W off +(1-α)*W on

[0066] Among them, W off and W on Respectively and The parameter weights, is the offline second image segmentation model, is the second online image segmentation model. α is the momentum coefficient, which is usually set to a large value to retain more pre-trained representations. By minimizing the downstream task loss and minimizing The output distribution difference is The parameter update L FT The formula is as follows:

[0067]

[0068] Among them, x l is the labeled data in the downstream task, L DS is the loss function of the downstream specific task. β is the loss coefficient used to adjust the impact of the two losses. By introducing It can not only solve downstream tasks specifically, but also alleviate the catastrophic forgetting of upstream representations, thereby achieving higher accuracy and stronger generalization.

[0069] Step S4: Adjust the parameters of the second image segmentation model through loss constraints and prediction constraints, and obtain a trained medical image processing model by combining the second image segmentation model after parameter adjustment with the first image segmentation model after self-supervised pre-training.

[0070] Specifically:

[0071] Build the upstream first image segmentation model and the downstream second image segmentation model network:

[0072] The basic structures of the upstream first image segmentation model and the downstream second image segmentation model are both PyramidTransformer U-Net (PTU). Specifically, the three upstream first image segmentation models are PTU-Large with 78M parameters, PTU-Hugev1 with 135M parameters, and PTU-Hugev2 with 131M parameters. The downstream second image segmentation model is PTU-Small with 49M parameters. (M, Million).

[0073] Verify the performance of the downstream second image segmentation model network after fine-tuning:

[0074] For the 3D medical image segmentation and detection dataset, the optimizer used for network training is the Adam with decoupled weight decay (AdamW) optimizer, and the loss function is a combination of Dice similarity coefficient loss and cross entropy loss. Since the encoder has been fully trained using the self-supervised learning method, this stage only requires fine-tuning the built 3D network using the downstream annotated dataset.

[0075] Based on the above method, the present invention provides a training system for a medical image processing model, including: an acquisition module, a training module, a migration module and a model generation module.

[0076] The acquisition module is used to obtain a number of unlabeled medical images and a number of labeled medical images, and to obtain a number of first image segmentation models and a second image segmentation model; wherein the complexity of the first image segmentation model is higher than that of the second image segmentation model; the training module is used to perform self-supervised pre-training on the first image segmentation models through the unlabeled medical images, obtain the self-supervised training representation with minimized prediction distribution, and transfer the self-supervised training representation to the second image segmentation model; the testing module is used to construct an online path and an offline path, and input the labeled medical images into the second image segmentation model in the online path and the offline path, and obtain the loss constraints of the downstream tasks and the prediction constraints of the offline path output; the model generation module is used to adjust the parameters of the second image segmentation model through the loss constraints and the prediction constraints, and obtain the trained medical image processing model by combining the second image segmentation model after parameter adjustment with the first image segmentation model after self-supervised pre-training.

[0077] The present invention also provides a computer device, including a memory and a processor, wherein a program is stored in the memory, and when the program is executed by the processor, the processor executes the steps of a training method for a medical image processing model.

[0078] In accordance with the disclosed embodiments, a computing device may communicate with one or more external devices (e.g., keyboards, pointing devices, Bluetooth communications, etc.), or with any device (e.g., routers, modems, etc.) that enables a computing device to communicate with one or more other computing devices.

[0079] The present invention also provides a storage medium on which a computer program is stored. When the computer program is executed by a processor, the steps of a training method for a medical image processing model are implemented.

[0080] According to the disclosed embodiments, the storage medium may be a non-volatile computer-readable storage medium, such as but not limited to: a portable computer disk, a hard disk, a random access memory (RAM), a read-only memory (ROM), an erasable programmable read-only memory (EPROM or flash memory), a portable compact disk read-only memory (CD-ROM), an optical storage device, a magnetic storage device, or any suitable combination thereof. In the present invention, the storage medium may be any tangible medium containing or storing a program that may be used by or in conjunction with an instruction execution system, apparatus, or device.

[0081] The above content is a further detailed description of the present invention in combination with a specific preferred embodiment. For technicians in the technical field to which the present invention belongs, several simple deductions or substitutions can be made without departing from the concept of the present invention, which should be regarded as belonging to the protection scope of the present invention.

Claims

1. A training method for a medical image processing model, characterized in that: The steps include: Acquire a number of unlabeled medical images and a number of labeled medical images, and acquire a number of first image segmentation models and a second image segmentation model; wherein the number of parameters of the first image segmentation model is higher than that of the second image segmentation model; Performing self-supervised pre-training on a plurality of first image segmentation models through the unlabeled medical images to obtain a self-supervised training representation with a minimized prediction distribution, and migrating the self-supervised training representation to a second image segmentation model; Constructing an online path and an offline path, and inputting the labeled medical image into a second image segmentation model in the online path and the offline path to obtain a loss constraint of a downstream task and a prediction constraint of an offline path output; The parameters of the second image segmentation model are adjusted through the loss constraint and the prediction constraint, and the trained medical image processing model is obtained by combining the second image segmentation model after parameter adjustment with the first image segmentation model after self-supervised pre-training.

2. A method for training a medical image processing model as claimed in claim 1, characterized in that: Using the unlabeled medical images, a plurality of first image segmentation models are self-supervised pre-trained, comprising the following steps: The standard teacher-student paradigm is used as self-supervised pre-training; the goal of self-supervised pre-training is to minimize the prediction distribution of the teacher model and the student model; During the self-supervised pre-training process, the parameters of the teacher model are updated by the exponential moving average of the student model parameters, and the parameters of the student model are updated by minimizing the output prediction of the teacher model to obtain a self-supervised training representation with minimized prediction distribution.

3. A method for training a medical image processing model as claimed in claim 2, characterized in that: Migrating the self-supervised training representation to a second image segmentation model comprises the following steps: In the first image segmentation model F (#,i) and the second image segmentation model F S Then add the same number of mappers to map the encoder output to the feature space; Given an unlabeled medical image training sample x, use the data augmentation module on the medical image training sample x Obtain the augmented image blocks x1 and x2; The x1 and x2 are respectively input into a pre-trained first image segmentation model and a randomly initialized second image segmentation model, and the features of the first image segmentation model and the second image segmentation model are aligned by introducing a consistency loss.

4. A method for training a medical image processing model as claimed in claim 3, characterized in that: The medical image training sample x is subjected to a data augmentation module. Obtain the augmented image blocks x1 and x2, the specific expression is: x1=τ1(x), x2=τ2(x) Among them, x1 and x2 are input into F (#，i) and F S middle.

5. A method for training a medical image processing model as claimed in claim 4, characterized in that: The method of aligning the features of the first image segmentation model and the second image segmentation model by introducing consistency loss comprises the following steps: The loss function in the consistency loss is: Among them, L mimic (·) is the loss function, They represent the first image segmentation models with different parameters pre-trained in three upstreams, e is an exponential function, t1 and t2 are temperature coefficients, σ ​​is a softMax function, Yes F S Mapper added later, P # Yes F (#，i) Mapper added later.

6. A method for training a medical image processing model as claimed in claim 3, characterized in that: The online path and offline path are constructed, and the specific expression is: W off =α*W off +(1-a)*W on Among them, W off and W on Respectively and The parameter weights, is the offline second image segmentation model, is the second online image segmentation model, The parameters are updated by The parameters are implemented as exponential moving average, and α is the momentum coefficient.

7. A method for training a medical image processing model as claimed in claim 6, characterized in that: Adjusting the parameters of the second image segmentation model by using the loss constraint and the prediction constraint comprises the following steps: For online paths By minimizing the downstream task loss and minimizing The output distribution difference is The parameter update L FT , the specific expression is: Among them, x l is the labeled data in the downstream task, L DS is the loss function of the downstream specific task, σ is the softMax function, and β is the loss coefficient used to adjust the impact of the two losses.

8. A training system for a medical image processing model, characterized in that: include: An acquisition module, used to acquire a number of unlabeled medical images and a number of labeled medical images, and acquire a number of first image segmentation models and a second image segmentation model; wherein the complexity of the first image segmentation model is higher than that of the second image segmentation model; A training module, configured to perform self-supervised pre-training on a plurality of first image segmentation models using the unlabeled medical images, obtain a self-supervised training representation with a minimized prediction distribution, and transfer the self-supervised training representation to a second image segmentation model; A testing module, used for constructing an online path and an offline path, and inputting the labeled medical image into a second image segmentation model in the online path and the offline path to obtain a loss constraint of a downstream task and a prediction constraint of an offline path output; The model generation module is used to adjust the parameters of the second image segmentation model through the loss constraint and the prediction constraint, and obtain the trained medical image processing model by combining the second image segmentation model after parameter adjustment with the first image segmentation model after self-supervised pre-training.

9. A computer device, characterized in that: It comprises a memory and a processor, wherein a program is stored in the memory, and when the program is executed by the processor, the processor executes the steps of a training method for a medical image processing model as described in any one of claims 1 to 7.

10. A storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the steps of a training method for a medical image processing model described in any one of claims 1 to 7 are implemented.

Images