Multi-stage loss optimization OCTA image classification method based on semi-supervised learning and knowledge distillation

Abstract

The invention discloses a multi-stage loss optimization OCTA (Optical Coherence Tomography Angiography) image diagnosis method based on semi-supervised learning and knowledge distillation. The method comprises the following steps: (1) constructing a teacher-student model framework in combination with semi-supervised learning; (2) data set processing: randomly removing part of sample labels of the training set as label-free data; (3) training a teacher model by using the labeled data, then generating a pseudo label of non-labeled data through the teacher model, and only retaining a high-quality pseudo label meeting a confidence threshold; and (4) training a student model by using a multi-stage loss function, wherein the student model comprises supervision loss of label data, supervision loss of pseudo-label data and knowledge distillation loss. Dependence on annotation data is reduced through semi-supervised learning, the performance of the lightweight student model is improved through knowledge distillation, and the method is suitable for scenes with limited clinical resources.

Description

Technical Field

[0001] The present invention belongs to the field of medical image diagnosis, and relates to a disease classification method for Optical Coherence Tomography Angiography (OCTA) images, specifically a multi-stage loss optimization method based on semi-supervised learning and knowledge distillation. Background Art

[0002] OCTA is a non-invasive medical imaging technology that is widely used in the detection and diagnosis of retinal and choroidal diseases. By capturing the details of the fundus microvascular network and tissue structure, OCTA provides an important basis for the early screening and diagnosis of diseases such as diabetic retinopathy and macular degeneration. However, the interpretation of OCTA images usually relies on experienced ophthalmologists, a process that is time-consuming and susceptible to human subjective factors.

[0003] Existing OCTA image classification methods are mainly based on deep learning technology, which usually requires a large amount of labeled data for supervised learning. However, the cost of obtaining high-quality medical labeled data is high, especially in clinical environments with limited resources, which has become a major obstacle to the large-scale application of deep learning technology. To solve this problem, semi-supervised learning methods have gradually attracted attention. By combining the learning of labeled data and unlabeled data, it significantly reduces the dependence on labeled data. Semi-supervised learning improves the utilization rate of unlabeled data by generating pseudo labels for unlabeled data using models trained with labeled data, but how to generate high-quality pseudo labels remains a key challenge.

[0004] While improving model performance, knowledge distillation methods are also widely used. Knowledge distillation is a model compression technology that transfers the knowledge of a high-performance complex model (teacher model) to a smaller lightweight model (student model), thereby significantly reducing computational overhead while ensuring model accuracy. This method provides an effective solution for model deployment in resource-constrained scenarios. However, existing knowledge distillation methods usually rely on labeled data for training and fail to make full use of unlabeled data.

[0005] The lesion information in OCTA images often has complex local features and global correlations, while traditional deep learning models are limited in capturing global dependencies and fine-grained features, especially when the disease types are complex and the differences are small, the classification performance is often not ideal. In addition, in order to adapt to the resource-constrained clinical environment, the research on lightweight deep learning models has gradually become a hot topic, but the performance of lightweight models is usually difficult to reach the level of complex deep models. How to combine semi-supervised learning to reduce the dependence on labeled data, and improve the performance of lightweight models through knowledge distillation, while effectively capturing local and global features in the image, has become an urgent problem to be solved in the current field of OCTA image classification. Summary of the invention

[0006] The purpose of the present invention is to propose a multi-stage loss optimization OCTA image classification method based on semi-supervised learning and knowledge distillation to address the problems in the prior art that OCTA image classification relies on a large amount of labeled data, the performance of lightweight models is insufficient, and the global feature capture capability is limited. This method constructs a semi-supervised learning framework, makes full use of labeled data and unlabeled data, generates high-quality pseudo labels, and reduces dependence on labeled data; at the same time, combined with knowledge distillation, the feature extraction capability of the complex model is transferred to the lightweight model, significantly improving the classification performance of the lightweight model. In addition, the present invention designs a multi-stage loss function, including supervision loss, pseudo-label supervision loss, and knowledge distillation loss, which effectively enhances the student model's ability to characterize local and global features, thereby achieving efficient and accurate classification and diagnosis of OCTA images, meeting application requirements in scenarios with limited clinical resources.

[0007] The technical solution for achieving the purpose of the present invention is:

[0008] A multi-stage loss optimization OCTA image classification method based on semi-supervised learning and knowledge distillation includes the following steps:

[0009] S1. Dataset division and processing

[0010] The OCTA-500 dataset is divided into a training set (105 labeled images and 350 unlabeled images), a validation set (15 labeled images), and a test set (20 labeled images). The training set is used for model training, the validation set is used to select the best teacher and student models, and the test set is used for final model performance evaluation. The unlabeled images in the training set will be used to generate pseudo labels.

[0011] S2. Semi-supervised learning - teacher model training and pseudo-label generation

[0012] The teacher model (ResNet50) is trained using 105 labeled data in the training set, and the cross entropy loss function is used to optimize the model. During the training process, the F1 score of the validation set is evaluated, and the teacher model weights with the best performance on the validation set are selected and used to generate pseudo labels. The trained teacher model is used to predict the 350 unlabeled images in the training set to generate pseudo labels. A confidence threshold (0.7) is set to retain only high-confidence pseudo labels for subsequent training, and samples below the threshold will be directly discarded, thereby improving the quality of pseudo labels and the robustness of training.

[0013] S3. Training of student model

[0014] The generated high-confidence pseudo labels are combined with 105 labeled data as input data sets for training the student model (MobileNetV2). During the training process, a multi-stage loss function is used for optimization, including:

[0015] Supervised loss: used for supervised learning with labeled data.

[0016] Pseudo-label supervised loss: For supervised learning with high-confidence pseudo-labeled data.

[0017] Knowledge distillation loss: Improves the performance of the student model by minimizing the difference between the prediction results of the student model and the teacher model.

[0018] A multi-stage loss function is used to optimize the student model.

[0019] S4. Testing and performance evaluation

[0020] The trained student model is applied to the test set, and the classification accuracy, F1 score, confusion matrix and other evaluation indicators of the model on different disease types (such as diabetic retinopathy, macular degeneration, etc.) are calculated to generate the final diagnosis results.

[0021] A multi-stage loss optimization OCTA image classification method based on semi-supervised learning and knowledge distillation, the system includes the following technical modules:

[0022] Data preprocessing module: responsible for the division of OCTA-500 data set and image preprocessing.

[0023] Teacher model training module: Use ResNet50 as the teacher model, train it on labeled data, and generate high-quality pseudo labels.

[0024] Pseudo-label generation module: predicts unlabeled data and only retains high-confidence pseudo-labels.

[0025] Student model training module: Use MobileNetV2 as the student model and perform optimization training through a multi-stage loss function.

[0026] Model evaluation module: Select the best model based on the F1 score of the validation set and evaluate the final performance on the test set.

[0027] Advantages or beneficial effects of this application:

[0028] This technical solution introduces a semi-supervised learning framework and uses unlabeled data to generate high-confidence pseudo-labels, significantly reducing dependence on labeled data and lowering data labeling costs. It is suitable for application scenarios with limited clinical resources, such as auxiliary diagnosis and screening of retinal diseases.

[0029] This technical solution combines knowledge distillation to transfer the complex feature extraction capabilities of the teacher model to the student model, while reducing the consumption of computing resources and enhancing the student model's ability to capture local and global features. It is suitable for clinical scenarios with limited resources.

[0030] This technical solution improves the classification ability of the lightweight student model by designing a multi-stage loss function, including supervision loss, pseudo-label supervision loss and knowledge distillation loss. While ensuring computational efficiency, it achieves high-precision diagnosis of OCTA image disease classification. BRIEF DESCRIPTION OF THE DRAWINGS

[0031] Figure 1 This is the overall structural flow chart of the present invention, showing the overall framework of data processing, semi-supervised learning and knowledge distillation.

[0032] Figure 2 This is a specific flow chart of semi-supervised learning in the present invention, showing the training process of the teacher model, the pseudo-label generation strategy, and the method for processing unlabeled samples.

[0033] Figure 3 This is a specific flow chart of knowledge distillation in this invention, showing the training process of the student model and the composition and optimization method of the multi-stage loss function. DETAILED DESCRIPTION

[0034] The present invention will be further described in detail below in conjunction with the accompanying drawings and specific embodiments, but the present invention is not limited thereto.

[0035] Example:

[0036] Reference Figure 1 , including the following steps:

[0037] S1. Dataset division and processing

[0038] 1.1) Dataset Overview: OCTA-500 contains annotated images of six major retinal disease categories, including age-related macular degeneration (AMD), diabetic retinopathy (DR), choroidal neovascularization (CNV), central serous chorioretinopathy (CSC), retinal vein occlusion (RVO) and other retinal diseases (Others).

[0039] 1.2) Data division: The data set is divided into training set, validation set and test set in a ratio of 7:1:2:

[0040] Training set: 105 labeled data + 350 unlabeled data (unlabeled data accounts for 75% of the training set);

[0041] Validation set: 15 labeled data;

[0042] Test set: 30 labeled data.

[0043] 1.3) Data preprocessing: The images are uniformly adjusted to a resolution of 512×512 and normalized to map pixel values ​​to the interval [0,1] to ensure that images from different sources have consistent input distribution.

[0044] S2. Semi-supervised learning - training of teacher model and pseudo-label generation, see Figure 2 , the specific steps are as follows:

[0045] 2.1) Teacher model training: ResNet50 is used as the teacher model and trained on the labeled data. The optimizer used for training is Adaptive moment estimation (AdamW); the initial learning rate is 0.001, and CosineDecay is used for dynamic adjustment; the weight decay coefficient is 0.0001 to suppress overfitting; batch size: 8; the maximum training rounds are 100, and the training is terminated early when the F1 score of the validation set has not improved for 20 consecutive rounds. During the training process, the best teacher model weight is selected by the F1 score of the validation set. The optimization target is the supervised loss function:

[0046]

[0047] in, is the teacher model's understanding of the input image The predicted output is is the corresponding true label, N L is the number of labeled data, and CE represents the cross-entropy loss function.

[0048] 2.2) Pseudo-label generation: Use the teacher model with the best F1 score on the validation set to predict the 350 unlabeled data in the training set and generate pseudo-labels. The pseudo-label generation rules are as follows:

[0049]

[0050] in is the confidence threshold function, For the teacher model to use unlabeled data The predicted probability of the sample is set to 0.7. If the predicted probability is less than 0.7, the sample will be directly discarded. The high-confidence pseudo-label generated is combined with the annotated data as the input data of the student model.

[0051] S3, student model training, refer to Figure 3 , the specific steps are as follows:

[0052] 3.1) Selection of student model: MobileNetV2 is used as the student model. This model is lightweight and suitable for resource-constrained clinical scenarios. The optimizer is Adam; the initial learning rate is 0.001, and Cosine Decay is used for dynamic adjustment; the weight decay coefficient is 0.0001; the batch size is 8; the maximum training rounds are 100, and the training is terminated early when the F1 score of the validation set does not improve for 20 consecutive rounds. During the training process, the best student model weight is selected by the F1 score of the validation set.

[0053] 3.2) Multi-stage loss function: The student model training process uses the following multi-stage loss function To optimize:

[0054]

[0055] in, is the supervision loss for labeled data; is the supervision loss for high confidence pseudo-labeled data; is the knowledge distillation loss; λ1=0.5, λ2=0.5 are the weights of the loss term.

[0056] The reasons for choosing the loss term weights λ1=0.5, λ2=0.5 are as follows:

[0057] (Labeled supervision loss) acts directly on the labeled data, can provide a strong supervision signal, and has higher credibility. Therefore, it is given a higher weight (1.0) to ensure that labeled data dominates model training. (Pseudo-label supervision loss) acts on unlabeled data, and its pseudo-label comes from the teacher model prediction. Although it increases the amount of training data, its label quality may be biased. Therefore, it is given a slightly lower weight (λ1 = 0.5) to avoid excessive interference with model training. (Knowledge distillation loss) is used to make the student model imitate the output of the teacher model and transfer the feature extraction ability of the teacher model. However, knowledge distillation loss may have a greater impact on the student model in the early stage of training, so a weight (λ2 = 0.5) is given to mitigate its effect to ensure that the student model can learn smoothly.

[0058] The detailed loss function is defined as follows:

[0059] Labeled supervised loss: Labeling data for the student model The predicted output, N L is the number of labeled data, and CE is the cross entropy loss function.

[0060] Pseudo-label supervision loss: N u is the number of pseudo-labeled data, is a pseudo label, Pseudo-labeled samples for the student model The predicted output.

[0061] Knowledge distillation loss: in is the prediction value of the teacher model for the unlabeled data, ∥·∥ 2 Represents the square of the Euclidean distance.

[0062] 3.3) The Adam optimizer back-propagates to update the student model weights. The process is as follows:

[0063]

[0064] m t is the first-order momentum estimate; v t is the first-order momentum estimate; β1, β2 are momentum coefficients (default β1 = 0.9, β2 = 0.999), is the gradient of the loss function; η is the learning rate (initial 0.001); ∈ is the smoothing parameter to prevent the denominator from being zero (∈=10 -8 ); γ is the weight attenuation coefficient (0.0001).

[0065] Dynamic adjustment of learning rate (Cosine Decay), by gradually reducing the learning rate during the training process, allows the model to converge quickly in the early stage and stabilize and optimize in the later stage, thereby avoiding oscillation and overfitting. This method helps the model find the optimal solution more accurately, while enhancing the robustness of training, and is suitable for long-term stable training tasks:

[0066]

[0067] η t is the current learning rate; η min is the minimum learning rate; η max is the initial learning rate (set to 0.001); t: current training round; T: total training rounds (set to 100).

[0068] S4. Model optimization and testing

[0069] The trained best student model is used to perform classification tests on the 20 labeled images in the test set. The model performance indicators include the following four types, which have different measurement functions and are used to comprehensively evaluate the classification performance. The higher the value, the better the corresponding model performance.

[0070] 4.1) Overall accuracy of multi-classification: The overall accuracy of multi-classification is the proportion of correctly predicted samples to the total number of samples;

[0071]

[0072] Wherein C is the total number of categories, that is, the number of label types involved in the classification task. In the task of the present invention, C=6.

[0073] TP i (True Positive for class i): The number of samples in category i that are correctly predicted as i.

[0074] FP i (False Positive for class i): The number of samples in category i that are incorrectly predicted to be i (actually belong to other categories).

[0075] FN i (False Negative for class i): The number of samples that actually belong to class i but are incorrectly predicted to belong to other categories.

[0076] 4.2) Multi-classification Precision: Multi-classification precision requires calculating the precision of each category and then taking the weighted average or macro average.

[0077] Precision per category:

[0078]

[0079] Macro-Precision:

[0080]

[0081] 4.3) Multi-classification Recall,Similarly, the recall rate can also be calculated for each category and then,the macro average is taken.

[0082] Recall per class:

[0083]

[0084] Macro-Recall:

[0085]

[0086] 4.4) The F1 score is the harmonic mean of precision and recall, and is used to measure the comprehensive performance of the classification model for precision and recall of each category. The F1 score can be calculated for each category separately and then the macro average is taken.

[0087] F1 score per category:

[0088]

[0089] Macro-F1:

[0090]

Claims

1. A multi-stage loss optimization OCTA image classification method based on semi-supervised learning and knowledge distillation, characterized in that: The following steps are involved: S1. Dataset division and preprocessing: The OCTA-500 dataset is divided into a training set (containing 105 labeled images and 350 unlabeled images), a validation set (15 labeled images), and a test set (20 labeled images) in a ratio of 7:1:

2. The images are uniformly adjusted to a resolution of 512×512 and normalized. S2. Training of the teacher model: ResNet50 is used as the teacher model, the labeled data of the training set is used to optimize the cross entropy loss function, and the best performing teacher model is selected based on the F1 score of the validation set; The optimization strategy is: Adaptive Moment Estimation (AdamW) optimizer is used, the initial learning rate is set to 0.001, and the learning rate is dynamically adjusted through Cosine Decay; The weight decay coefficient is set to 0.0001; the training batch size is set to 8; the maximum number of training rounds is 100, and the training is terminated early when the F1 score of the validation set does not improve for 20 consecutive rounds; S3. Pseudo-label generation: Use the trained teacher model to predict the unlabeled data and generate pseudo-labels based on the set confidence threshold (0.7). Unlabeled data below the threshold will be discarded. 3.1) Predict unlabeled data using the confidence threshold function: Screening high-confidence samples; 3.2) Use the unlabeled data with confidence greater than 0.7 as pseudo labels and combine them with the labeled data to train the student model; S4. Student model training: MobileNetV2 is used as the student model. The generated high-confidence pseudo labels are combined with labeled data. The student model is optimized through a multi-stage loss function, where the multi-stage loss function includes the supervision loss of labeled data, the supervision loss of high-confidence pseudo labels, and the knowledge distillation loss. 4.1) Adaptive Moment Estimation (AdamW) optimizer was used, with the initial learning rate set to 0.001; the weight decay coefficient was set to 0.0001; the training batch size was set to 8; the maximum number of training rounds was 100, and training was terminated early when the F1 score of the validation set did not improve for 20 consecutive rounds; 4.2) The multi-stage loss function is as follows: in, Supervised loss for labeled data (cross entropy loss); Supervision loss (cross entropy loss) for high-confidence pseudo-labeled data; Knowledge distillation loss (Euclidean distance); λ1=0.5, λ2=0.5 are the weights of the loss terms; Labeled supervised loss: CE is the cross entropy loss function; Pseudo-label supervision loss: is a pseudo label; Knowledge distillation loss: in is the predicted value of the teacher model for the unlabeled data. ∥·∥ 2 represents the square of the Euclidean distance; S5. Model optimization and verification: The student model is optimized through the dynamic learning rate adjustment strategy (Cosine Decay), and the best student model weight is selected according to the F1 score of the validation set; S6. Testing and performance evaluation: Apply the trained student model to the test set, output the classification results, and perform performance evaluation through multi-class classification accuracy, macro-average F1 score, macro-average recall, and macro-average precision.

Images