A training method, system, device and medium for a medical image processing model

Abstract

The present invention discloses a training method, system, device and medium for a medical image processing model, which relates to the field of model training technology and includes the following steps: obtaining a number of unlabeled medical images and constructing a number of first image segmentation models and a second image segmentation model; inputting the unlabeled medical images into the number of first image segmentation models for self-supervised pre-training; migrating the self-supervised training representation of the first image segmentation model to the second image segmentation model; constructing an online path and an offline path in the migrated second image segmentation model, and the second image segmentation model in the online path will be used for the final downstream task test. The present invention proposes an adaptation method for self-supervised pre-training and downstream task fine-tuning, which can enable the second image segmentation model to learn the self-supervised representation of the first image segmentation model while alleviating the catastrophic forgetting problem of the self-supervised representation.

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 tremendous success in 3D image analysis. To ensure high accuracy and generalizability of trained models, a large amount of labeled data is generally required. However, in medical image analysis, labeling 3D images is extremely time-consuming and labor-intensive, and requires the expertise 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, enabling the learning of visual representations without manual annotation. For example, the paper "Ibot: Image Bert Pre-training with Online Tokenizer," published at the 10th International Conference on Learning Representations, is a self-supervised learning method that ensures that the input mask image remains consistent within the constraints of the teacher-student model. The feature extractor trained using this method has demonstrated excellent performance in downstream tasks such as image segmentation, classification, and detection.

[0004] However, current self-supervised learning methods typically fine-tune the upstream pre-trained first image segmentation model directly on small amounts of downstream data. For downstream tasks with small amounts of labeled data, a second image segmentation model is often more appropriate. Furthermore, when fine-tuning the first image segmentation model on downstream tasks, catastrophic forgetting of the upstream pre-trained model representation often occurs. This means that the fine-tuned first image segmentation model gradually overfits to the downstream task, significantly reducing the accuracy and generalization of the fine-tuned model. 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 in the prior art 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.

[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 plurality of unlabeled medical images and labeled medical images, and construct a plurality 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 using the unlabeled medical images to obtain self-supervised training representations with minimized prediction distributions, and migrating the self-supervised training representations 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 for a downstream task and a prediction constraint for 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 parameter-adjusted second image segmentation model with the self-supervised pre-trained first image segmentation model.

[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 a minimized prediction distribution.

[0015] Preferably, migrating the self-supervised training representation into 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 Get 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 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 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 upstream models, e is the exponential function, t1 and t2 are the temperature coefficients, σ ​​is the softMax function, It's F S Mapper added later, P # It's F (#,i) Mappers added later.

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

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

[0028] Among them, W off and W on Respectively and The parameter weights of 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 of 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, configured to acquire a plurality of unlabeled medical images and a plurality of labeled medical images, and to acquire a plurality 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 of a plurality of first image segmentation models using the unlabeled medical images to obtain self-supervised training representations with minimized prediction distributions, and to transfer the self-supervised training representations to a second image segmentation model;

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

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

[0038] The present invention provides a computer device comprising 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, enabling the second image segmentation model to have the representation ability of the first image segmentation model. When fine-tuning the second image segmentation model on downstream small data, an online path and an offline path are constructed. 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, 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 This is the overall structural 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 accompanying drawings. Obviously, the described embodiments are part of the embodiments of the present invention, not all of them. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative efforts 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 pre-trained parameters are fixed. Given an unlabeled training sample x, the data augmentation module is used on it The augmented image patches x1 and x2 are obtained and fed into the pre-trained first image segmentation model and the randomly initialized second image segmentation model, respectively. A consistency loss is introduced to align the features of the first and second image segmentation models, enabling the transfer of the pre-trained representations of the first image segmentation model to the second image segmentation model. For downstream small-data tasks, the parameters of the second image segmentation model in the online path are fine-tuned using the loss constraints of the downstream task and the prediction constraints on 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 testing.

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

[0046] Step S1: Acquire a number of unlabeled medical images and a number of labeled medical images, and construct a number of first image segmentation models and a second image segmentation model, wherein the complexity and number of parameters 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 for self-supervised pre-training. 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 a minimized prediction distribution.

[0051] Step S2: self-supervised pre-training of several first image segmentation models using unlabeled medical images to obtain self-supervised training representations with minimized prediction distribution, and migrating the self-supervised training representations 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, Represents three upstream pre-trained first image segmentation models with different parameter values. 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 capabilities of the first image segmentation model. It's F S Mapper added later, P # It's F (#,i) Mappers 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 for fine-tuning the downstream task. The two paths initially have the same model weights. and in, The parameters are updated by The parameters are implemented as exponential moving average. The online path and offline path are constructed. The specific expressions are:

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

[0066] Among them, W off and W on Respectively and The parameter weights of is the offline second image segmentation model, is the second image segmentation model in the online path. α 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 of 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 not only solves downstream tasks specifically, but also alleviates 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 parameter-adjusted second image segmentation model with the self-supervised pre-trained first image segmentation model.

[0070] Specifically:

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

[0072] The basic architecture of both the upstream first image segmentation model and the downstream second image segmentation model is a 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.

[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 network training used the Adam with decoupled weight decay (AdamW) optimizer, and the loss function was a combination of Dice similarity loss and cross-entropy loss. Since the encoder has been fully trained using self-supervised learning methods, this stage only requires fine-tuning the built 3D network using downstream annotated datasets.

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

[0076] The acquisition module is used to acquire a number of unlabeled medical images and a number of labeled medical images, and to 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; the training module is used to perform self-supervised pre-training on the number of first image segmentation models through unlabeled medical images, obtain a 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 to 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 parameter-adjusted second image segmentation model with the self-supervised pre-trained first image segmentation model.

[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] According to the disclosed embodiments, a computing device may communicate with one or more external devices (e.g., a keyboard, a pointing device, Bluetooth communications, etc.), or with any device that enables a computing device to communicate with one or more other computing devices (e.g., a router, a modem, etc.).

[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 can 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 specific preferred embodiments. For those skilled in the art to which the present invention belongs, several simple deductions or replacements can be made without departing from the concept of the present invention, which should be regarded as falling within the scope of protection of the present invention.

Claims

1. A training method for a medical image processing model, characterized in that: The steps include: Acquire a plurality of unlabeled medical images and a plurality of labeled medical images, and acquire a plurality 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; performing self-supervised pre-training on a plurality of first image segmentation models using the unlabeled medical images to obtain self-supervised training representations with minimized prediction distributions, and migrating the self-supervised training representations 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 for a downstream task and a prediction constraint for an offline path output; Adjusting the parameters of the second image segmentation model using the loss constraint and the prediction constraint, and combining the parameter-adjusted second image segmentation model with the self-supervised pre-trained first image segmentation model to obtain a trained medical image processing model; Migrating the self-supervised training representation to a second image segmentation model includes the following steps: In the first image segmentation model and the second image segmentation model Then add the same number of mappers to map the encoder output to the feature space; Given an unlabeled medical image training sample x , for medical image training samples x Using the Data Augmentation Module , get the augmented image block and ; Respectively and Input into the pre-trained first image segmentation model and the randomly initialized second image segmentation model, and align the features of the first image segmentation model and the second image segmentation model by introducing consistency loss; The method of aligning the features of the first image segmentation model and the second image segmentation model by introducing consistency loss includes the following steps: The loss function in the consistency loss is: in, (·) is the loss function, , respectively represent the first image segmentation models with different parameter amounts pre-trained in three upstreams, is an exponential function, and is the temperature coefficient, is the softMax function, yes The mapper added later, yes Mappers added later.

2. A training method for a medical image processing model according to claim 1, characterized in that: Performing self-supervised pre-training on a plurality of first image segmentation models using the unlabeled medical images comprises 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 a minimized prediction distribution.

3. The training method of a medical image processing model according to claim 1, wherein: The medical image training samples x Using the Data Augmentation Module , get the augmented image block and , the specific expression is: in, and Enter separately and middle.

4. The training method of a medical image processing model according to claim 1, wherein: The online path and offline path are constructed as follows: in, and Respectively and The parameter weights of 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. is the momentum coefficient.

5. The training method of a medical image processing model according to claim 4, wherein: Adjusting the parameters of the second image segmentation model using the loss constraint and the prediction constraint includes the following steps: For online paths , by minimizing the downstream task loss and minimizing The output distribution difference of Parameter update , the specific expression is: in, is the labeled data in the downstream task, is the loss function of the downstream specific task, is the softMax function, is the loss coefficient used to adjust the effects of the two losses.

6. A system for training a medical image processing model according to any one of claims 1 to 5, characterized in that: include: an acquisition module, configured to acquire a plurality of unlabeled medical images and a plurality of labeled medical images, and to acquire a plurality 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 of a plurality of first image segmentation models using the unlabeled medical images to obtain self-supervised training representations with minimized prediction distributions, and to transfer the self-supervised training representations to a second image segmentation model; A testing module is configured to construct an online path and an offline path, and input the labeled medical image into a second image segmentation model in the online path and the offline path to obtain a loss constraint for a downstream task and a prediction constraint for an offline path output; A model generation module is used to adjust the parameters of the second image segmentation model through the loss constraint and the prediction constraint, and to obtain a trained medical image processing model by combining the parameter-adjusted second image segmentation model with the self-supervised pre-trained first image segmentation model.

7. A computer device, characterized in that: It includes 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 training method of a medical image processing model as described in any one of claims 1 to 5.

8. 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 a medical image processing model according to any one of claims 1 to 5 are implemented.

Images