Medical image classification model parameter equalization fine tuning method based on target gradient equilibrium

Abstract

The invention discloses a medical image classification model parameter equalization fine tuning method based on target gradient equilibrium, which belongs to the field of transfer learning of medical image classification, and comprises the following steps of: dynamically adjusting gradient amplitudes of parameters of different layers in a fine tuning process by constructing a target gradient equalization module; and the overfitting problem caused by scarcity of medical image data is effectively relieved. Firstly, parameters which meet high correlation and are distributed in all layers are screened, and the selected parameters are used for sparse fine tuning. And then, projection of the gradient of the head sample in the gradient direction of the tail sample is utilized to assist the tail sample to complete updating, so that the overall fine adjustment effect is improved. And finally, re-weighting the parameters, and carrying out balance fine tuning on the model based on the re-weighting loss according to contributions of different categories to the total loss.

Description

Technical Field

[0001] This invention belongs to the field of transfer learning in medical image classification, specifically for fine-tuning imbalanced medical image datasets. The product is mainly used for fine-tuning medical image classification. Background Technology

[0002] Traditional deep learning-based medical image classification methods often suffer from unstable performance due to insufficient training samples and rare cases. A common approach is transfer learning, particularly fine-tuning, which involves pre-training on a large source dataset and then fine-tuning on a smaller target dataset to improve training results. However, when classifying rare cases with limited sample sizes, the model's gradients and weights tend to favor head samples (those with more data) during fine-tuning. This hinders optimal fine-tuning on imbalanced medical datasets, causing the classification model to fail to identify tail cases (those with fewer data) or to make inaccurate judgments, thus affecting diagnostic accuracy and the safety of automated diagnosis. Therefore, we have developed a fine-tuning method that effectively addresses the impact of insufficient target dataset size and class imbalance on classification in medical image transfer learning.

[0003] On the other hand, existing deep learning-based medical image classification methods face two bottlenecks that urgently need to be addressed: insufficient sample size and imbalanced data classification. Common data-related techniques include AutoAugment, Fast AutoAugment, RandAugment, and RankMix. Resampling techniques, such as SMOTE, Borderline-SMOTE, R-SMOTE, and KMeansSMOTE, are also relevant. While these methods can mitigate the impact of class imbalance to some extent, oversampling can lead to overfitting, and undersampling can result in the loss of valuable information, limiting the improvement of model performance. In terms of model algorithms, a widely used method is class balancing loss. Although this method helps alleviate class imbalance when combined with data resampling, it typically relies on passive adjustment based on the loss function, which may reduce stability and lacks targeted training of model parameters. From a hybrid approach perspective, a novel multi-weighted re-loss function and an end-to-end cumulative learning strategy have been proposed to address the problems of imbalanced sample size and high classification difficulty. However, these still require sufficient training data within the same domain, and their applicability may be limited when transferring the model between different domains.

[0004] Furthermore, in the medical image classification task, three main fine-tuning strategies were employed: selective layer fine-tuning, regularization, and model adaptation. Selective layer fine-tuning includes selective unfreezing and selective feature extraction. Selective unfreezing methods include linear probing, unfreezing from front to back or back to front, or combinations of various unfreezing strategies, allowing for the selective extraction of features from different layers of the pre-trained model as needed. Regularization-based fine-tuning strategies, similar to commonly used regularization methods, are based on L1 and L2 regularization. Finally, model adaptation strategies involve structural modifications to the transfer model, such as deleting or adding layers, to improve the model's adaptability to the target domain task. Despite these significant advancements, excessive parameters in deep models can lead to overfitting during training, especially when data is limited. Additionally, class imbalance can hinder optimal fine-tuning on the target dataset. Summary of the Invention

[0005] Purpose of the invention: This invention aims to provide a parameter balancing and fine-tuning method for medical image classification models based on target gradient balancing, in order to solve the overfitting problem that may occur in the parameters of the target domain and the pre-trained model.

[0006] Technical solution: A method for fine-tuning parameters of a medical image classification model based on target gradient balancing, the method comprising: Parameters that are highly correlated and distributed across all layers are selected and used for sparse fine-tuning to address potential overfitting issues in the target domain and pre-trained model parameters. By projecting the gradient of the head sample onto the gradient direction of the tail sample, the tail sample can be updated, thereby improving the overall fine-tuning effect. The model parameters are reweighted based on the contribution of different categories to the overall loss, enabling the model to undergo a fine-tuning based on the reweighted loss to mitigate the bias caused by class imbalance.

[0007] Furthermore, the specific implementation process of this method includes: S1. Sparse fine-tuning based on gradient attribution: Obtain the pre-trained CNN model and the target domain medical image dataset; calculate the gradient values ​​of each parameter in the pre-trained CNN model under the target domain medical image classification task; the gradient values ​​are obtained by taking the partial derivative of the loss function with respect to the parameters, the loss function being cross-entropy loss; for each neuron in the pre-trained CNN model, select the top A parameters with the largest gradient values ​​from all the parameters corresponding to the input connections of that neuron; all selected parameters together form the parameter group; then perform sparse fine-tuning on the selected parameter group to complete the initial parameter update of the pre-trained CNN model under the target domain medical image classification task; The parameters determined by the screening are distributed across all layers of the pre-trained CNN model; A is a preset positive integer whose value range is dynamically adjusted according to the model structure and the complexity of the target task. S2. Gradient projection-based auxiliary fine-tuning: During the sparse fine-tuning process in step S1, the gradient components corresponding to the head samples and the tail samples in the target domain medical image dataset are obtained respectively; then, the directional relationship between the gradient components of the head samples and the gradient components of the tail samples is determined. If there is no conflict between the two, the gradient components of the head sample are projected onto the gradient components of the tail sample, and the projection results are used to assist the tail sample in completing the parameter update. If there is a conflict between the two, the cosine similarity loss function is used to optimize the gradient components of the head sample and the tail sample, and then the optimized gradient components are used to assist the model in completing the parameter update. S3. Based on loss reweighting, the overall loss function of the pre-trained CNN model is set as cross-entropy loss. The number M of categories participating in training in the target domain medical image dataset is determined, and an initial contribution is set for each category. According to the contribution of each category to the overall loss, the loss weight corresponding to that category is calculated, and the loss weight is inversely proportional to the category contribution. The cross-entropy loss is reweighted and updated using the calculated loss weights to obtain the updated loss function. Based on the updated loss function, the model after the auxiliary fine-tuning in step S2 is balanced and fine-tuned to complete the final parameter update of the model in the target domain medical image classification task, and obtain the fine-tuned model for medical image classification.

[0008] Furthermore, the gradient values ​​of the parameters in the model are calculated using the following expressions:

[0009] The parameters of the pre-trained CNN model are represented as follows: , It is the cross-entropy loss function.

[0010] Furthermore, in step S2, and Let these be the gradient components of the head sample and the tail sample, respectively. There is no conflict between their gradient components, and they are represented as follows: Then, the gradient components of the head samples are projected onto the direction of the gradient components of the tail samples, and the gradient component update of the tail samples is expressed as follows: : ; If the gradient components of the head sample and the tail sample conflict, when using the cosine similarity loss function to optimize the gradient components of the head sample and the tail sample, the calculation formula of the cosine similarity loss function is as follows: .

[0011] Step S3 involves reweighting the parameters and updating the model based on the contribution of samples from different categories to the loss, setting the overall loss function of the model as cross-entropy loss, and finally determining the loss. The update is represented as:

[0012] In the formula, This indicates the number of categories participating in the training. This represents the initial contribution of each category. This indicates the contribution of each category to the overall loss; The updated calculation of the loss weight ratio is as follows:

[0013] The loss of the medical classification model based on target gradient balancing is... Finally, the final classification result is obtained by continuously optimizing and fine-tuning the model.

[0014] Furthermore, the pre-trained CNN model includes the ResNet50 model or other deep convolutional neural network models used for image classification.

[0015] On the other hand, the present invention provides a system for medical image classification, comprising a memory and a processor. The memory stores a computer program, and when the processor executes the computer program, it implements the steps of the parameter equalization fine-tuning method for medical image classification described above. The system is used for classifying medical images. The medical images include at least one of CT images, MRI images, and ultrasound images. The system includes: The gradient attribution-based sparse fine-tuning module executes step S1 of the parameter equalization fine-tuning method for medical image classification in the method described above. The gradient projection-based auxiliary fine-tuning module executes step S2 of the parameter equalization fine-tuning method for medical image classification in the method described above. Based on the loss reweighting equalization fine-tuning module, step S3 of the parameter equalization fine-tuning method for medical image classification is executed.

[0016] On the other hand, the present invention also provides a computer-readable storage medium storing a computer program that, when executed by a processor, implements the steps of the parameter equalization fine-tuning method for medical image classification.

[0017] Beneficial effects: This invention proposes a parameter balancing fine-tuning method for medical classification models based on target gradient balancing, which can alleviate the problem of insufficient target dataset affecting classification results to a certain extent, and solve the problem of imbalance in target dataset categories affecting the final classification effect, effectively improving the fine-tuning effect of medical image classification. Attached Figure Description

[0018] Figure 1 This is an implementation architecture diagram of the method described in this invention; Figure 2 This is a schematic diagram illustrating how the gradient of the head sample assists in updating the gradient of the tail sample. Detailed Implementation

[0019] In this invention, we propose a parameter balancing fine-tuning method for medical classification models based on target gradient balancing, applicable to medical image classification. It consists of a sparse fine-tuning module based on gradient attribution, an auxiliary fine-tuning module based on gradient projection, and a balancing fine-tuning module based on loss reweighting. First, parameters that satisfy high relevance and distribution across all layers are selected. These selected parameters are used for sparse fine-tuning to address potential overfitting issues in the target domain and pre-trained model parameters. Then, due to the significant difference in sample size between the head and tail data, the gradient of the head samples can be projected onto the gradient direction of the tail samples to assist in updating the tail samples, thereby improving the overall fine-tuning effect. Finally, because of the imbalance in data size between the head and tail, model weights often favor the head samples. Parameters can be reweighted, and the model can undergo balancing fine-tuning based on the contribution of different categories to the overall loss, using reweighted loss.

[0020] Specifically, combined Figure 1 The implementation steps of this invention are as follows: S1. Sparse Fine-tuning Based on Gradient Attribution First, during the pre-training process of a CNN model, different parameters extract different features, and the same parameters contribute differently to fine-tuning across various tasks. Specifically, given a target domain dataset and a pre-trained CNN model... The ultimate goal is to find a set of parameters. , Perform efficient fine-tuning.

[0021] This parameter set needs to meet two conditions: it must be highly relevant to the classification task of the target domain and be distributed across all layers.

[0022] Related research indicates that the parameter with the largest gradient value represents the parameter whose loss function changes the fastest and has the largest rate of change along the gradient direction of that parameter, and is beneficial for effective gradient descent during fine-tuning. Therefore, the importance of parameters in a pre-trained model for a specific task can be identified by selecting the parameter with the highest gradient value. The gradient of a parameter is calculated as follows:

[0023] in, It is a specific loss function; this invention chooses the cross-entropy loss function. Because... Equivalent to Therefore, selecting the parameter group with the largest change in gradient value is equivalent to selecting the parameter group with the largest change in parameter value during fine-tuning.

[0024] Furthermore, when sparsely fine-tuning the pre-trained model, the selected parameter sets should be distributed throughout the entire network. For each neuron in the network, the top-A parameters are selected from all input connections of that neuron. In this way, the fine-tuned network will not only focus on the semantic feature information of the gradients of deep neurons, but also pay attention to some primary features extracted by shallow neurons that are different from the source domain.

[0025] S2, Gradient Projection-Based Auxiliary Fine-Tuning

[0026] First, during the fine-tuning process, the gradient components of the head and tail samples are obtained separately. and Different projection operations are performed to address the differences in gradient directions between head and tail samples. For conflict-free gradient components... ,like Figure 2 Then, the gradient components of the head samples are projected onto the direction of the gradient components of the tail samples, and the gradient component update of the tail samples is expressed as follows: :

[0027] for Figure 2 (b) shows the conflicting gradient components. Furthermore, to avoid the gradient direction at the head and the gradient direction at the tail appearing Figure 2 (c) shows the complete conflict phenomenon ( =-1), based on the cosine similarity loss function (CosineEmbeddingLoss) to optimize the gradients of the head and tail.

[0028] Gradient components of the head sample gradient components of tail samples The cosine similarity loss function between them can be calculated using the following formula: .

[0029] Finally, the optimized gradient-assisted model is used for fine-tuning.

[0030] S3. Equilibrium Fine-tuning Based on Loss Reweighting

[0031] In addition, due to the difference in the amount of data between the head samples and the tail samples, the model weights will be biased towards the head samples. Therefore, the parameters can be re-weighted and the model updated by combining the contribution of samples of different categories to the loss.

[0032] First, the overall loss function of the model is set to cross-entropy loss (CE loss). The number of classes participating in training is set to M, and the initial contribution of each class is... Then the final loss Updated to: .

[0033] When fine-tuning the model, the loss value corresponding to each category is selected from the loss of each batch. The inverse proportion of category contribution is used as a weight to re-sum the losses, and the updated CE loss is used for model updates. The updated calculation of the loss weight proportions is as follows:

[0034] Therefore, the overall loss of the GBF model is Finally, the final classification result is obtained by continuously optimizing and fine-tuning the model.

[0035] Table 1. Comparison results with mainstream methods in ResNet50

Claims

1. A method for fine-tuning parameters of a medical image classification model based on target gradient balancing, characterized in that, The method includes: Parameters that are highly correlated and distributed across all layers are selected and used for sparse fine-tuning to address potential overfitting issues in the target domain and pre-trained model parameters. By projecting the gradient of the head sample onto the gradient direction of the tail sample, the tail sample can be updated, thereby improving the overall fine-tuning effect. The model parameters are reweighted based on the contribution of different categories to the overall loss, enabling the model to undergo a fine-tuning based on the reweighted loss to mitigate the bias caused by class imbalance.

2. The method for fine-tuning the parameter balancing of a medical image classification model according to claim 1, characterized in that, The specific implementation process of this method includes: S1. Sparse fine-tuning based on gradient attribution: Obtain the pre-trained CNN model and the target domain medical image dataset; calculate the gradient values ​​of each parameter in the pre-trained CNN model under the target domain medical image classification task; the gradient values ​​are obtained by taking the partial derivative of the loss function with respect to the parameters, the loss function being cross-entropy loss; for each neuron in the pre-trained CNN model, select the top A parameters with the largest gradient values ​​from all the parameters corresponding to the input connections of that neuron; all selected parameters together form the parameter group; then perform sparse fine-tuning on the selected parameter group to complete the initial parameter update of the pre-trained CNN model under the target domain medical image classification task; The parameters determined by the screening are distributed across all layers of the pre-trained CNN model; A is a preset positive integer whose value range is dynamically adjusted according to the model structure and the complexity of the target task. S2. Gradient projection-based auxiliary fine-tuning: During the sparse fine-tuning process in step S1, the gradient components corresponding to the head samples and the tail samples in the target domain medical image dataset are obtained respectively; then, the directional relationship between the gradient components of the head samples and the gradient components of the tail samples is determined. If there is no conflict between the two, the gradient components of the head sample are projected onto the gradient components of the tail sample, and the projection results are used to assist the tail sample in completing the parameter update. If there is a conflict between the two, the cosine similarity loss function is used to optimize the gradient components of the head sample and the tail sample, and then the optimized gradient components are used to assist the model in completing the parameter update. S3. Based on loss reweighting, the overall loss function of the pre-trained CNN model is set as cross-entropy loss. The number M of categories participating in training in the target domain medical image dataset is determined, and an initial contribution is set for each category. According to the contribution of each category to the overall loss, the loss weight corresponding to that category is calculated, and the loss weight is inversely proportional to the category contribution. The cross-entropy loss is reweighted and updated using the calculated loss weights to obtain the updated loss function. Based on the updated loss function, the model after the auxiliary fine-tuning in step S2 is balanced and fine-tuned to complete the final parameter update of the model in the target domain medical image classification task, and obtain the fine-tuned model for medical image classification.

3. The method for fine-tuning the parameter balancing of a medical classification model according to claim 1 or 2, characterized in that, The gradient values ​​of the parameters in the model are calculated using the following expressions: The parameters of the pre-trained CNN model are represented as follows: , It is the cross-entropy loss function.

4. The method for fine-tuning the parameter balancing of a medical classification model according to claim 1 or 2, characterized in that, In step S2, and Let these be the gradient components of the head sample and the tail sample, respectively. There is no conflict between their gradient components, and they are represented as follows: Then, the gradient components of the head samples are projected onto the direction of the gradient components of the tail samples, and the gradient component update of the tail samples is expressed as follows: : ； If the gradient components of the head sample and the tail sample conflict, when using the cosine similarity loss function to optimize the gradient components of the head sample and the tail sample, the calculation formula of the cosine similarity loss function is as follows: 。 5. The method for fine-tuning the parameter balancing of a medical classification model according to claim 1 or 2, characterized in that, Step S3 involves reweighting the parameters and updating the model based on the contribution of samples from different categories to the loss, setting the overall loss function of the model as cross-entropy loss, and finally determining the loss. The update is represented as: In the formula, This indicates the number of categories participating in the training. This represents the initial contribution of each category. This indicates the contribution of each category to the overall loss; The updated calculation of the loss weight ratio is as follows: The loss of the medical classification model based on target gradient balancing is... Finally, the final classification result is obtained by continuously optimizing and fine-tuning the model.

6. The method for fine-tuning the parameter balancing of a medical classification model according to claim 1 or 2, characterized in that, The pre-trained CNN model includes the ResNet50 model or other deep convolutional neural network models used for image classification.

7. A system for medical image classification, characterized in that, The system includes a memory and a processor. The memory stores a computer program. When the processor executes the computer program, it implements the steps of the parameter equalization fine-tuning method for medical image classification as described in any one of claims 1-6. The system is used for classifying medical images.

8. The system for medical image classification according to claim 7, characterized in that, The medical images include at least one of CT images, MRI images, and ultrasound images; The system includes: The gradient attribution-based sparse fine-tuning module executes step S1 of the parameter equalization fine-tuning method for medical image classification in the method described above. The gradient projection-based auxiliary fine-tuning module executes step S2 of the parameter equalization fine-tuning method for medical image classification in the method described above. Based on the loss reweighting equalization fine-tuning module, step S3 of the parameter equalization fine-tuning method for medical image classification is executed.

9. A computer-readable storage medium, characterized in that, The computer-readable storage medium stores a computer program that, when executed by a processor, implements the steps of the parameter equalization fine-tuning method for medical image classification as described in any one of claims 1-6.

Images