A medical image classification method based on federated learning and contrastive learning

Abstract

The application discloses a medical image classification method based on federated learning and contrast learning, in the training of federated learning, the local model of the current training round and the global model of the last round are combined with the local model for training, the distance between the model of the current round and the global model of the last round is reduced, the distance between the model of the current round and the local model of the last round is increased, the supervised contrast loss between the models is used, and the local model of the client gradually approaches the global model in the training process. This method can effectively reduce the performance decline problem caused by data heterogeneity and increase the generalization ability of the model.

Description

Technical Field

[0001] This invention relates to the field of image processing, and specifically to a medical image classification method, which is based on federated learning and contrastive learning. Background Technology

[0002] Image classification is one of the most common and fundamental tasks in computer vision. Essentially, it's an image processing method that distinguishes images from different categories based on the different features expressed by semantic information within the image. Examples include image classification and recognition of skin diseases such as melanoma, benign melanocytic lesions, basal cell carcinoma, and actinic keratosis, or detection of chest X-ray images of COVID-19 and other bacterial infections. However, due to data privacy concerns and the heterogeneity of data from different institutions, current medical image classification methods are not entirely effective.

[0003] Data heterogeneity refers to the fact that data from different medical institutions are not independently and identically distributed. Due to different sampling methods, the datasets of different institutions often vary greatly, which can lead to a severe decline in model performance. Federated learning has made significant contributions to addressing data privacy issues; however, when training on data with severe heterogeneity, it can cause serious differences between the models trained by the clients, resulting in a significant performance degradation as the number of communication rounds increases. Summary of the Invention

[0004] The purpose of this invention is to provide a medical image classification method based on federated learning and contrastive learning. This method can effectively alleviate the performance degradation caused by data heterogeneity and increase the generalization ability of the model.

[0005] To address the problems existing in the prior art, the technical solution adopted in this invention is as follows:

[0006] A medical image classification method based on federated learning and contrastive learning is proposed. In the training of federated learning, the local model of the current training round and the global model of the previous round are combined with the local model for training. This reduces the distance between the feature representations of the current model and the global model of the previous round, and increases the distance between the feature representations of the current model and the local model of the previous round. By using supervised contrastive loss between models, the local model of the client gradually approaches the global model during the training process, thereby improving the performance of the model and mitigating the impact of data heterogeneity.

[0007] Advantages or effects of the present invention:

[0008] Federated learning allows medical institutions to conduct cross-institutional collaborative training while storing data locally, effectively solving the data privacy protection problem. However, in practical applications of medical image classification, the heterogeneity of data among different clients leads to the final trained classification model failing to learn good feature representations. To address this issue, this invention proposes a federated contrastive learning method to mitigate the performance degradation caused by data heterogeneity. It combines local and global models for contrastive learning, enabling each local model to learn semantic information from other client data, thereby promoting better feature representations and improving the model's generalization ability.

[0009] To demonstrate the effectiveness of the method, experiments were conducted on two datasets used for medical image classification, and good experimental results were obtained in both cases. This shows that, without considering communication costs, the method exhibits strong robustness in heterogeneous data and has good application prospects. Attached Figure Description

[0010] Figure 1 This is a framework diagram for a medical image classification method based on federated learning and contrastive learning.

[0011] Figure 2 To address different numbers of clients, a performance curve comparison chart of the medical image classification method based on federated learning and contrastive learning in this invention with other federated learning methods is presented in a skin lesion dataset.

[0012] Figure 3 To address different numbers of clients, this invention presents a performance curve comparison chart of its medical image classification method based on federated learning and contrastive learning with other federated learning methods on a COVID-19 test dataset. Detailed Implementation

[0013] The present invention will be further described below with reference to the embodiments and accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0014] The following will provide a detailed explanation from three aspects: background and core technologies, implementation methods, comparative experiments, and ablation experiments. Background and Core Technologies:

[0015] 1. Federated Learning

[0016] Federated learning is a distributed machine learning technique that allows each collaborating institution (such as a hospital) to conduct cross-institutional collaborative training by storing data locally. During training, only model parameters are exchanged, without exchanging local data, thus ensuring collaborative model training without the need for data sharing and achieving data privacy protection.

[0017] The training process for federated learning mainly consists of the following steps:

[0018] (1) The central server sends the current global model to each client; the current model comes from the global model after the (t-1)th round of aggregation. It is distributed to each client as a local model for local training. Let the local model of the k-th (k = 1, 2, ..., K) local client be denoted as .

[0019] (2) Client S k The model is trained using local data; each client, after receiving the model from the central server, trains it using a cross-entropy loss function and stochastic gradient descent. Indicates client S k The loss function used for training is given by equation (1), where η represents the learning rate during model training.

[0020]

[0021] (3) Each client sends the trained model to the central server; after training is complete, client S k Send its local model to the server;

[0022] (4) The central server aggregates local models; the server is responsible for storing the models. Aggregation is achieved using a weighted average method, typically with weights represented by the size of the data set. This relates to the global model. The aggregation is shown in formula (2):

[0023]

[0024] Where |D| represents the total number of samples in each client dataset, |D k | Indicates that it comes from client S k The number of samples of local data.

[0025] The goal of federated learning is to learn a global model w from various clients without exchanging data, with the help of a central server. Therefore, the overall optimization objective of the model is as shown in formula (3):

[0026]

[0027] Among them, L k (w) represents the client S k The cross-entropy loss during model training is given by the formula (4):

[0028]

[0029] In formula (4), D k Indicates client S k In the corresponding dataset, x represents a sample in the data, and y represents a label in the data.

[0030] 2. Comparative learning

[0031] Contrastive learning, as a method within self-supervised learning, attempts to learn good feature representations from unlabeled data. Simultaneously, in the field of supervised learning, supervised contrastive learning methods have also emerged. These methods effectively utilize class label information to generate a feature space for image data of different classes. Images belonging to the same class are grouped together in the feature space, while images belonging to different classes are separated. This approach significantly improves the accuracy of image classification tasks.

[0032] This invention primarily studies supervised contrastive learning, which is derived from self-supervised learning. For labeled data, the data can be divided into positive and negative pairs based on the class label; that is, data from the same class label form a positive pair, and data from different class labels form a negative pair. The goal of supervised contrastive learning is to reduce the distance between feature representations from images of the same class and increase the distance between feature representations from images of different classes.

[0033] set up Representing the image x i All images facing each other, similarly, for image x i The generated feature vector z i Its positive and negative pairs are and The definition of its supervised contrastive loss function is shown in formula (5):

[0034]

[0035] in Indicate z i The number of positive samples, where τ represents the temperature parameter.

[0036] Example:

[0037] like Figure 1As shown, a medical image classification method based on federated learning and contrastive learning is presented. The trained model framework consists of a Densely Connected Convolutional Network (DenseNet) backbone for classification, a linear mapping layer, and a fully connected layer. DenseNet is used to extract image features, the linear mapping layer is used to map the feature vectors to a specified dimension, and the fully connected layer is used to generate the predicted value for each class. The training for the t-th round of communication includes the following steps:

[0038] Step 1: The server sends the global model for the current round to each client, and the client processes the global model to train the local model for the current round.

[0039] Step 2: The client stores the local model to be trained, the global model distributed by the server, and the local model trained in the previous round. It then uses a hybrid loss function based on supervised contrastive loss and cross-entropy loss for training. The supervised contrastive loss function is shown in formula (6):

[0040]

[0041] Where sim() represents the cosine similarity function, and τ is a temperature parameter used to control the loss value and ensure convergence of the loss. This represents the feature representation learned by the local model in round t. The feature representation learned by the local model in round t-1 is represented by . The feature representation learned by the global model in the (t-1)th round;

[0042] The formula for the hybrid loss function used in training is shown in (7):

[0043]

[0044] Where s is the local dataset sample currently being trained. For the local model to be trained in the current t-th round, For the local model in round t-1, Let μ be the global model for the (t-1)th round, and μ be a hyperparameter used to control the weights of the supervised comparison loss.

[0045] Step 3: After each client finishes training locally, it sends the trained local model to the server for model aggregation.

[0046] Step 4: After collecting all the local models from the clients, the server aggregates the models using a weighted average to obtain the global model for round t. At this point, round t of communication training is complete; the server then returns to step 1 for the next round of communication training. The overall optimization objective is shown in formula (8):

[0047]

[0048] Where n represents the total amount of data from all clients, n k This represents the number of datasets from the k-th client. For the mixed loss of the k-th client, w glob Since this is the global model to be trained, the overall optimization objective of the model is to use a mixture loss for each client. After training, a weighted average is then used to aggregate the results, ultimately yielding the optimal global model w. glob .

[0049] Comparative experiments and ablation experiments:

[0050] 1. Experimental Preparation

[0051] Using the DenseNet classification network as the backbone, the dataset needed to be partitioned to simulate local clients in order to simulate the experimental setup of federated learning. To simulate data heterogeneity caused by class imbalance, the Dirichlet distribution Dir(β) was used to generate data for each local client. With this strategy, clients may have fewer or no data samples for some class labels. This experiment used two datasets: ISIC 2018: Task 1, a dataset for classifying skin lesions of melanoma; and Task 2, a dataset for COVID-19 testing. The evaluation metrics for the experimental results were Accuracy, Specificity, Sensitivity, AUC, and F1 score.

[0052] 2. Comparative Experiment

[0053] In the experiments, the method of this invention was compared with FedAvg, FedProx, and centralized training methods. The experimental results on the two datasets are shown in Tables 1 and 2, respectively, where the data in bold represents the data of the method of this invention.

[0054] Table 1 compares our method with current state-of-the-art federated learning methods in Task 1.

[0055] Method Accuracy Specificity Sensitivity AUC F1 Central 88.87％ 87.27％ 74.11％ 91.05％ 73.48％ FedAvg 87.92％ 85.66％ 75.55％ 90.55％ 73.86％ FedProx 88.59％ 86.29％ 74.08％ 91.52％ 73.61％ FedCL(ours) 89.24％ 87.92％ 77.30％ 91.92％ 75.00％

[0056] Table 2 compares our method with current state-of-the-art federated learning methods in Task 2.

[0057] Method Accuracy Specificity Sensitivity AUC F1 Central 92.99％ 94.32％ 88.98％ 96.85％ 88.42％ FedAvg 92.36％ 94.19％ 89.25％ 96.47％ 87.66％ FedProx 92.84％ 94.11％ 89.81％ 97.07％ 88.39％ FedCL(ours) 93.43％ 94.77％ 90.90％ 97.31％ 89.38％

[0058] Tables 1 and 2 show that the FedAvg and FedProx methods achieved good results on both datasets. However, in most cases, the performance of federated learning cannot surpass that of the method without federated learning. Notably, compared to these methods, the federated contrastive learning method of this invention achieved the best results on all classification metrics, exceeding them by approximately 1% on some evaluation metrics. This demonstrates that this method can effectively address the problem of data heterogeneity, reduce its impact on model performance, and achieve better generalization.

[0059] 3. Ablation test

[0060] (1) The impact of different weights on experimental results

[0061] This experiment investigated different weights for the mixed loss, and the results are shown in Table 3. It can be seen that when the weight is 1.1 or 0.9, the performance of most evaluation metrics decreases, indicating that optimal performance can only be obtained by selecting appropriate weights. When the weight is set to 1.0, the performance of most evaluation metrics is balanced, resulting in the best experimental results.

[0062] Table 3 shows the comparative experiments under different weights for Tasks 1 and 2.

[0063]

[0064] (2) The impact of different numbers of clients on model performance

[0065] This experiment was conducted with different numbers of clients using three different methods. The experimental results are as follows: Figure 2 and Figure 3As shown in the figure, AUC is used as the evaluation metric for the experiment. In Task 1, the number of clients is set to 8, 9, 10, 11, and 12, while in Task 2, the number of clients is set to 3, 4, 5, 6, and 7. The line graph shows that the generalization performance of federated learning decreases with the increase of the number of clients. This is because, compared to centralized training, data dispersion makes it more difficult for the model to learn good feature representations, leading to a decrease in classification performance. However, the method of this invention, FedCL, achieves the best performance under different client number settings. More importantly, with the increase of the number of local clients, the method of this invention degrades more slowly than other methods. In Task 1, when the number of local clients is 8, the method of this invention outperforms FedAvg by only 0.26%, while when the number of local clients is 12, the method of this invention outperforms FedAvg by 1.75%. This indicates that other federated learning methods suffer significant performance degradation due to data heterogeneity, while the FedCL method of this invention significantly reduces the impact of data heterogeneity.

[0066] Although embodiments of the invention have been shown and described, it will be understood by those skilled in the art that various changes, modifications, substitutions and alterations can be made to these embodiments without departing from the principles and spirit of the invention, the scope of which is defined by the appended claims and their equivalents.

Claims

1. A medical image classification method based on federated learning and contrastive learning, characterized in that, In federated learning training, the local model of the current training round and the global model of the previous round are combined with the local model for training. This reduces the distance between the feature representations of the current model and the global model of the previous round, and increases the distance between the feature representations of the current model and the local model of the previous round. The supervised contrast loss between models is used to allow the client's local model to gradually approach the global model during the training process. The trained model framework consists of a backbone network for classification, a densely connected convolutional network (DenseNet), a linear mapping layer, and a fully connected layer. DenseNet extracts image features, the linear mapping layer maps the feature vectors to a specified dimension, and the fully connected layer generates the predicted value for each class. The communication training for each round includes the following steps: Step 1: The server sends the global model for the current round to each client, and the client processes the global model to train the local model for the current round. Step 2: The client stores the local model to be trained, the global model distributed by the server, and the local model trained in the previous round. It uses a hybrid loss function based on supervised contrastive loss and cross-entropy loss for training. The supervised contrastive loss function is shown in formula (6): （6） in, Represents the cosine similarity function. Temperature parameters are used to control the loss value and ensure loss convergence. Representative of the first The feature representation learned by the local model in the round. Representing the The feature representation learned by the local model in the round. Representing the The feature representation learned by the global model in the round; The formula for the hybrid loss function used in training is shown in (7): （7） in, For the current training local dataset samples, For the current number Local models awaiting training. For the first The local model of the wheel, For the first The global model of the wheel This is a hyperparameter used to control the weights of the supervised comparison loss; Step 3: After each client finishes training locally, it sends the trained local model to the server for model aggregation. Step 4: After collecting all the local models from all clients, the server aggregates the models using a weighted average method to obtain the [model name missing]. The global model of the wheel, up to this point, the first... The round of communication training is complete; the process returns to step 1 for the next round of communication training; the overall optimization objective is shown in formula (8): （8） in, Represents the total amount of data across all clients. This represents the number of datasets from the k-th client. For the mixed loss of the k-th client, Since this is the global model to be trained, the overall optimization objective of the model is to use a mixed loss for each client. After training, a weighted average is then used to aggregate the results, ultimately yielding the optimal global model. .

Images