A cardiac arrhythmia classification method based on contrastive learning

Abstract

The present invention discloses a method for classifying arrhythmias based on contrastive learning, and relates to the technical field of electrocardiogram classification. The present invention comprises at least the following steps: inputting an electrocardiogram signal, performing two different data enhancements on the electrocardiogram signal, and thereby constructing a positive and negative sample pair, wherein the positive sample is a new electrocardiogram signal obtained after the same electrocardiogram signal has been subjected to different data enhancements. The present invention imitates the data enhancement method of an image, and adopts two different wavelet transform (ratio = 0.5) + segmented random occlusion methods for the electrocardiogram signal. In the contrastive learning model, a suitable data enhancement method is crucial. The method adopted by the present invention has a positive effect on the contrast of the positive sample pairs, and obtains the lower branch by pruning the upper branch to find the long-tail data, and implicitly increases the weight through the loss function. A dynamic temperature coefficient strategy is adopted for the loss function. The method adopted by the present invention improves the balance and enhances the generalization and migration capabilities of the model.

Description

Technical Field

[0001] The present invention relates to the technical field of electrocardiogram classification, and in particular to an arrhythmia classification method based on contrast learning. Background Art

[0002] An electrocardiogram (ECG) records the heart's depolarization and repolarization processes during each cardiac cycle. It reflects the heart's electrophysiological processes and helps doctors diagnose arrhythmias. Holter monitoring, an important tool for early screening and postoperative follow-up of arrhythmias, generates a large number of single-lead ECGs, posing a significant challenge to the healthcare system. Because the morphologies of various arrhythmias vary slightly, and abnormal waveforms do not always appear continuously, it is difficult for doctors alone to analyze these massive records. Therefore, intelligent analysis of ECG signals has attracted widespread attention from academia and industry both domestically and internationally.

[0003] The key to supervised deep learning methods is a large, labeled ECG dataset. However, ECG data involves patient privacy concerns, and ECG annotation requires the expertise of physicians, making it prohibitively expensive. Several studies have been proposed to mitigate the negative impact of data scarcity.

[0004] Data augmentation is an effective method to expand training samples. Pan et al. proposed and tested several one-dimensional data augmentation methods, proving that data augmentation can improve the performance of deep learning models in detecting arrhythmias, but the improvement effect is limited. Hatamian et al. used Generative Adversarial Networks (GAN) to generate atrial fibrillation type ECG data to expand samples. The experimental results show that the classification performance on the atrial fibrillation category can be improved, but in some scenarios, the detection performance of the normal type will be impaired, which shows that the performance of the data augmentation method will be affected by the scale and category distribution of the original data. Hui Chen et al. proposed CLECG, applying the contrastive learning model to the ECG field to mine information from unlabeled data.

[0005] However, the effectiveness of CLECG on real outpatient electrocardiograms has not been verified. Compared with the organized public datasets, real outpatient electrocardiograms are often more unbalanced in categories and have a long-tail distribution. Moreover, due to unknown labels, they cannot be balanced through sampling, making them difficult to be well applied in real-world scenarios. Summary of the Invention

[0006] The purpose of the present invention is to provide an arrhythmia classification method based on contrastive learning to solve the existing problem: the effect of CLECG on real outpatient electrocardiograms has not been verified. Compared with the sorted public data sets, real outpatient electrocardiograms are often more unbalanced in categories and have a long-tail distribution. Moreover, due to the unknown labels, they cannot be balanced through sampling, and are difficult to be well applied in real scenarios.

[0007] To achieve the above object, the present invention provides the following technical solution: a method for classifying arrhythmias based on contrastive learning, comprising at least the following steps:

[0008] S1: Input an ECG signal and perform two different data augmentations on it to construct positive and negative sample pairs. The positive sample is a new ECG signal obtained by augmenting the same ECG signal with different data augmentations, and the negative sample is an ECG signal obtained by augmenting other ECG signals with data augmentation.

[0009] S2: Prune the upper branch to obtain the lower branch, thereby creating a self-competitive model, where the upper branch is the target model and the lower branch is the self-competitive model;

[0010] S3: Sample the input signal and apply two different data augmentations to create two different versions [V i , V' i ];

[0011] S4: [V i , V' i ]by[f i (·), f'1(·)] encoding, and their output features [f1, f2] are fed into the nonlinear projection head to force similarity under the NT-Xent loss. There are at least two cases, one of which is ideal, which is able to remember the sample well, then pruning will not "forget" it, so it will not cause additional perturbations. In the other case, for rare and atypical cases, the prediction difference between the pruned model and the unpruned model will be amplified, so the weight of these samples will be implicitly increased in the overall loss;

[0012] S5: Negative samples are saved using a memory queue, which increases the number of negative samples.

[0013] S6: Through the loss function, the distance between positive samples is shortened and the distance between negative samples is increased;

[0014] S7: In the loss function, a dynamic temperature coefficient is used;

[0015] S8: Finally, in fine-tuning, load the one-dimensional encoder as the initialization parameters of the training target dataset and randomly initialize the classifier to adapt to the target dataset.

[0016] Preferably, two different data enhancements are used for the ECG signal in S1, including at least two different wavelet transforms (ratio = 0.5) + segmented random occlusion. The segmented operation can allow a larger occlusion ratio to ensure the difference between positive samples, while avoiding large sections of continuous signals being set to zero, resulting in training collapse.

[0017] Preferably, the pruning adopts amplitude-based pruning, and the pruning ratio is set to 0.2.

[0018] Preferably, the pruning mask of f'1(·) is updated in each iteration after the model weights are updated, and the application includes at least the following steps:

[0019] The backbone network is a large DNN whose weights do not change much in one or two iterations, so I set the pruning mask to be updated lazily at the beginning of each epoch to save computational overhead;

[0020] All iterations in the same period use the same mask;

[0021] Since the self-competitor is always derived and updated from the latest target model, the two branches will co-evolve during training.

[0022] Preferably, when updating two branches, [f i (·), f'1(·)] will share the same weights in the unpruned part, and f i (·) will update the rest independently, and the weights corresponding to f'1(·) will be pruned to zero.

[0023] Preferably, the loss function is expressed as:

[0024]

[0025] where Z i and z' i Represents the same sample x i Projection representation of different data transformations;

[0026] τ represents the temperature coefficient, which is used to adjust the difficulty of the task. The smaller the temperature coefficient, the more attention is paid to the most difficult negative samples to distinguish.

[0027] A dynamic temperature coefficient approach is adopted, setting upper and lower bounds and an adjustment period for the temperature coefficient. Long-tail data benefits from a smaller temperature coefficient, while non-long-tail data benefits from a larger temperature coefficient. This is used to continuously improve the separation between classes in long-tail data without any additional computational cost.

[0028] Preferably, the S5 at least includes the following steps:

[0029] A memory queue is used to solve the balance problem between the number of negative samples and the batch size. A queue is used to save the sample expressions in several adjacent past batches.

[0030] The saved small queues are different from the images in the current batch, so they are all considered negative samples, thus decoupling the number of negative samples from the batch size;

[0031] Due to the existence of the queue, the key branch cannot calculate the gradient normally. Therefore, a momentum update mechanism is introduced to use the parameters of the query branch to perform momentum updates on the parameters of the key branch.

[0032] Compared with the prior art, the present invention has the following beneficial effects:

[0033] 1: This paper imitates the data enhancement method of images and adopts two different wavelet transform (scale = 0.5) + segmented random occlusion methods for ECG signals. In the contrastive learning model, a suitable data enhancement method is crucial. The method adopted by this paper has a positive effect on the comparison of positive sample pairs.

[0034] 2. The present invention uses pruning to create a self-competitive model. By enforcing similarity between the self-competitive model and the target model, the weight of the long-tail data is implicitly increased, and the imbalance of the ECG data is reduced, which is very meaningful for real-world applications.

[0035] 3. This paper establishes a memory bank of size K to store samples from several adjacent batches in the key branch. Due to the characteristics of neural network training, these samples do not overlap with the samples in the current batch and can therefore be considered negative samples for the current batch. The introduction of the bank effectively expands the number of negative samples while maintaining a relatively small batch size.

[0036] 4. The present invention adopts a dynamic temperature coefficient strategy in the loss function, sets appropriate upper and lower bounds of the temperature coefficient and the adjustment period, and takes advantage of the characteristics that long-tail data benefits from a smaller temperature coefficient and non-long-tail data benefits from a larger temperature coefficient to perform continuous task switching, thereby improving balance without introducing any additional costs. BRIEF DESCRIPTION OF THE DRAWINGS

[0037] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the following briefly introduces the drawings required for describing the embodiments. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without creative work.

[0038] Figure 1 This is the effect diagram of the data enhancement method adopted by the present invention;

[0039] Figure 2 A flow chart of comparative learning adopted by the present invention;

[0040] Figure 3 It is a fine-tuning flow chart of the present invention. DETAILED DESCRIPTION

[0041] The technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, rather than all the embodiments.

[0042] See also Figure 1-3 As shown, the present invention is a method for arrhythmia classification based on contrastive learning, comprising the following steps:

[0043] Input an ECG signal I;

[0044] By performing two different data enhancements, the data enhancement method is two different wavelet transforms (scale = 0.5) + segmented random occlusion;

[0045] Here, two different wavelet transforms are used in this embodiment: Daubechies 4 wavelet transform (db4) and Daubechies 6 wavelet transform (db6);

[0046] In addition, considering that ECG signals are usually long (e.g., 10,000 sample points), we propose to randomly crop them in segments;

[0047] More specifically, the signal is first divided into k equal segments, each segment containing N signal points.

[0048] Then, in each segment, the number of consecutive p×N signal points is randomly set to O, where P is the ratio of random cropping;

[0049] We get two data-enhanced ECG signals [V i , V' i ],like Figure 1 As shown;

[0050] These two ECG signals are obtained by using different data augmentation methods on the same ECG signal. We call them positive sample pairs, and other ECG data are negative sample pairs.

[0051] Each data must be processed in this way, that is, N ECG data will be obtained after data enhancement.

[0052] Among them, different transformed perspectives from the same sample point constitute a positive sample pair, while conversely, the transformed perspectives of different sample points constitute a negative sample pair.

[0053] like Figure 2 As shown, we have a one-dimensional encoder f i (·) performs amplitude-based pruning to obtain the one-dimensional encoder f'1(·).

[0054] We set the pruning rate to 20%.

[0055] The upper branch is called the target model, and the pruned lower branch is called the self-competition model.

[0056] ECG signal [V i , V' i ]After one-dimensional encoder [f i (·), f'1(·)] encoding

[0057] Extract features to get features [f1, f2]

[0058] Then send the output feature [f1, f2] to the mapping head [g i (·), g'1(·)],

[0059] Enforce similarity in the loss function.

[0060] In each iteration, we will use the simplest amplitude-based pruning. Ideally, the pruning mask for f'1(·) can be updated at each iteration after the model weights are updated. In practice, since the backbone network is a large DNN and its weights do not change much in one or two iterations, we set the pruning mask to be updated lazily at the beginning of each epoch to save computational overhead;

[0061] All iterations in the same epoch adopt the same mask. Since the self-competitor is always obtained and updated from the latest target model, the two branches will co-evolve during the training process.

[0062] When we update two branches, [f i (·), f'1(·)] will share the same weights in the unpruned part, and f i (·) will update the rest independently (corresponding to the weights of f'1(·) being pruned to zero).

[0063] If f i (·) remembers the sample well, then prune f i (·) It will not be “forgotten” and thus will not cause additional perturbations.

[0064] Otherwise, the prediction differences between the pruned and unpruned models will be amplified for rare and atypical cases, so the weights of these samples will be implicitly increased in the overall loss.

[0065] The loss function is:

[0066]

[0067] where Z i and z' i Represents a positive sample pair, z' i and Z k represents a negative sample pair.

[0068] The meaning of this loss function is to shorten the distance between positive samples and increase the distance between negative samples, so as to learn the potential features of the data in pre-training and improve the arrhythmia classification performance.

[0069] In the loss function, τ refers to the temperature coefficient. Generally speaking, the temperature coefficient is a fixed coefficient in the loss function.

[0070] In the present invention, we set a dynamic temperature coefficient and a corresponding adjustment period by taking advantage of the fact that long-tail data benefits from a smaller temperature coefficient and non-long-tail data benefits from a larger temperature coefficient.

[0071] Set a suitable upper and lower bounds for the temperature coefficient and a suitable adjustment period. We set up an experiment to find the most suitable upper and lower bounds for the temperature coefficient and the adjustment period.

[0072] As shown in Table 1 and Table 2:

[0073] Table 1

[0074] 0.03 0.015 0.01 0.005 0.05 0.689 0.693 0.690 0.682 0.005 0.708 0.716 0.712 0.707 0.001 0.697 0.711 0.708 0.701

[0075] Table 2

[0076] T No 100 300 500 1000 0.688 0.709 0.716 0.710 0.692

[0077] From the experimental results, we can see that the performance is best when the upper limit of the temperature coefficient is 0.015, the lower limit is 0.05, and the adjustment cycle is 300.

[0078] We set the upper and lower bounds of the temperature coefficient to (0.005, 0.015) and the period T to T = 300

[0079] We also conducted experiments to compare the performance of the pre-training of this invention with the CLECG model. The data set used was PTB-XL, and the indicator was ACC.

[0080] After undersampling the PTB-XL dataset, we compared our proposed SDCLECG method with the traditional CLECG method.

[0081] Comparison of linear separability performance between CLECG and SDCLECG

[0082] Framework Many Medium Few All CLECG 0.627 0.604 0.550 0.583 SDCLECG 0.653 0.635 0.618 0.621

[0083] Comparing the few-shot performance of CLECG and SDCLECG

[0084] Framework Many Medium Few All CLECG 0.532 0.518 0.482 0.493 SDCLECG 0.568 0.537 0.504 0.529

[0085] Table 3

[0086] Framework Acc CLECG 0.685 Ours 0.716

[0087] It can be seen that the ACC score of the present invention reaches 0.716, which exceeds the score of CLECG.

[0088] After pre-training, we obtain the trained one-dimensional encoder i(·).

[0089] like Figure 3 As shown in Figure 2, in the fine-tuning stage, the one-dimensional encoder i(·) is loaded as the initialization parameter of the training target dataset and the classifier is randomly initialized to adapt to the target dataset.

[0090] We conducted experiments on the generalization and transfer capabilities of the fine-tuning stage performance.

[0091] Table 4: Migration ability test results

[0092] ICBEB2018 data volume AUC 1 / 8 0.9364 1 / 4 0.9587 1 / 2 0.9646 1 0.9792

[0093] From the experimental results, it can be seen that the present invention has achieved good AUC scores regardless of whether it is one-eighth of the data volume or the full data volume, indicating that the pre-trained model obtained by the present invention has the ability to extract the common features of ECG signals, and the knowledge obtained can be better transferred to different target tasks.

[0094] Table 5: Generalization ability experimental results

[0095] PhysioNet2017 Randomly initialize F1 Pre-trained F1 Resetnet1d_wang 0.6510 0.6814 ResNet501d 0.6207 0.6535 Inception1d 0.6449 0.6613 ICBEB2018 Randomly initialized AUC Pre-training AUC Resetnet1d_wang 0.8890 0.9221 ResNet501d 0.9071 0.9378 Inception1d 0.9122 0.9413

[0096] The dataset PhysioNet 2017 is an atrial fibrillation dataset. As can be seen from the table, both the arrhythmia classification task and

[0097] Atrial fibrillation detection tasks can benefit from the pre-trained model trained on PTB-XL.

[0098] The performance on varies with the network structure.

[0099] It will be apparent to those skilled in the art that the present invention is not limited to the details of the exemplary embodiments described above and that the invention can be embodied in other specific forms without departing from the spirit or essential characteristics of the invention. Therefore, the embodiments should be considered in all respects as illustrative and non-restrictive, and the scope of the invention is defined by the appended claims, not the foregoing description, and all variations within the meaning and range of equivalents of the claims are intended to be included therein. Any reference sign in a claim should not be construed as limiting the claim to which it relates.

Claims

1. A method for classifying arrhythmias based on contrastive learning, characterized by: At least the following steps are included: S1: Input an ECG signal and perform two different data augmentations on it to construct positive and negative sample pairs. The positive sample is a new ECG signal obtained by augmenting the same ECG signal with different data augmentations, and the negative sample is an ECG signal obtained by augmenting other ECG signals with data augmentation. S2: Prune the upper branch to obtain the lower branch, thereby creating a self-competitive model, where the upper branch is the target model and the lower branch is the self-competitive model; S3: Sample the input signal and apply two different data augmentations to create two different versions [V i , V' i ]; S4: [V i , V' i ]by[f i (·), f'1(·)] encoding, and their output features [f1,f2] are fed into the nonlinear projection head to enforce similarity under the NT-Xent loss. There are at least two cases, one of which is that if the sample is well remembered, pruning will not "forget" it and thus will not cause additional perturbations. In the other case, the prediction difference between the pruned model and the unpruned model will be amplified, so the weights of these samples will be implicitly increased in the overall loss. S5: Negative samples are saved using a memory queue, which increases the number of negative samples. S6: Through the loss function, the distance between positive samples is shortened and the distance between negative samples is increased; S7: In the loss function, a dynamic temperature coefficient is used; S8: Finally, in fine-tuning, load the one-dimensional encoder as the initialization parameters of the training target dataset and randomly initialize the classifier to adapt to the target dataset.

2. The arrhythmia classification method based on contrastive learning according to claim 1, characterized in that: In S1, two different data enhancements are used for ECG signals, including at least two different wavelet transforms + segmented random occlusion.

3. The arrhythmia classification method based on contrastive learning according to claim 1, characterized in that: The pruning adopts amplitude-based pruning, and the pruning ratio is set to 0.

2.

4. The arrhythmia classification method based on contrastive learning according to claim 1, characterized in that: The pruning mask of f'1(·) is updated at each iteration after the model weights are updated, and the application includes at least the following steps: The backbone network is a large DNN whose weights do not change much in one or two iterations, so the pruning mask is set to be updated lazily at the beginning of each epoch to save computational overhead; All iterations in the same period use the same mask; Since the self-competitor is always derived and updated from the latest target model, the two branches will co-evolve during training.

5. The arrhythmia classification method based on contrastive learning according to claim 1, characterized in that: When updating two branches, [f i (·), f'1(·)] will share the same weights in the unpruned part, and f i (·) will update the rest independently, and the weights corresponding to f'1(·) will be pruned to zero.

6. The arrhythmia classification method based on contrastive learning according to claim 1, characterized in that: The loss function is expressed as: where Z i and z' i Represents the same sample x i Projection representation of different data transformations; τ represents the temperature coefficient, which is used to adjust the difficulty of the task. The smaller the temperature coefficient, the more attention is paid to the most difficult negative samples to distinguish. A dynamic temperature coefficient approach is adopted, setting upper and lower bounds and an adjustment period for the temperature coefficient. Long-tail data benefits from a smaller temperature coefficient, while non-long-tail data benefits from a larger temperature coefficient. This is used to continuously improve the separation between classes in long-tail data without any additional computational cost.

7. The arrhythmia classification method based on contrastive learning according to claim 1, characterized in that: The S5 at least includes the following steps: A memory queue is used to solve the balance problem between the number of negative samples and the batch size. A queue is used to save the sample expressions in several adjacent past batches. The saved small queues are different from the images in the current batch, so they are all considered negative samples, thus decoupling the number of negative samples from the batch size; Due to the existence of the queue, the key branch cannot calculate the gradient normally. Therefore, a momentum update mechanism is introduced to use the parameters of the query branch to perform momentum updates on the parameters of the key branch.

Images