Diversity active learning method based on generative adversarial network

Abstract

The application discloses a diversity active learning method based on a generative adversarial network, and realizes landing application in a railway train fault detection related data set in combination with a classification task model ResNet-18. The strategy comprises the following steps: preparing a railway train fault image classification data set, selecting a classification task model, training a generative adversarial network by using the data set and generating a synthetic sample; using a task model feature map to perform feature extraction and matching of the synthetic sample and the original sample; performing sample screening according to an adaptive selection decision based on a network parameter gradient norm; finally, introducing a K-means++ clustering algorithm as a diversity strategy to perform sample screening; and repeatedly iterating the above sample generation and screening steps to realize active learning sample selection. The strategy improves the active learning decision performance by combining the above multi-dimensional strategies, and provides effective help for reducing project annotation cost.

Description

Technical Field

[0001] This invention belongs to the field of active learning in machine learning, and particularly relates to a diversity-based active learning method based on generative adversarial networks. Background Technology

[0002] Active learning is a research area within machine learning. Specifically, active learning is a strategy that aims to filter samples, selecting high-value samples to provide for training supervised learning models. By reducing the annotation costs of supervised learning models, it achieves training results that approach or even surpass those of fully supervised learning models.

[0003] Among these methods, the starting point is to use existing samples to synthesize new samples as labeled samples through a generator. With the explosive research boom in Generative Adversarial Networks (GANs) in recent years, using generators in GANs for sample synthesis and combining it with active learning has become a research hotspot.

[0004] Currently, research combining active learning with Generative Adversarial Networks (GANs) is relatively limited. GAAL was the first study to combine these two approaches, but its performance based purely on GAN-based active learning was unsatisfactory. Furthermore, due to the performance limitations of early GANs, it was unable to generate high-quality images at medium to high resolutions. The subsequent ASAL approach addressed GAN's insufficient generative performance by indirectly utilizing synthesized samples from GANs and labeling similar samples within existing datasets, achieving some performance improvement, but it did not directly use generated images to train the model. While both GAAL and ASAL strategies directly or indirectly utilize GANs for active learning, their final active learning performance was unsatisfactory due to flawed strategies and the limitations of GAN performance.

[0005] Active learning, combined with supervised learning models, has a wide range of applications, including image classification, object detection, and human pose estimation. In image classification, it is also used in various industrial production scenarios and biomedical fields such as CT scans, gastroscopy images, and polyp detection. Since annotating medical images presents a high barrier to entry for non-medical personnel, introducing active learning can significantly reduce annotation costs. Therefore, this strategy is also applied to railway train fault detection in industrial production, where it can also significantly reduce annotation costs. Summary of the Invention

[0006] The purpose of this invention is to provide a diversity-based active learning method based on generative adversarial networks (GANs) to address the technical problem that existing active learning strategies based on GANs do not perform well in image multi-classification tasks.

[0007] To solve the above-mentioned technical problems, the specific technical solution of the present invention is as follows:

[0008] A diversity-based active learning method based on generative adversarial networks includes the following steps:

[0009] Step S1: Prepare an unlabeled image classification dataset and use it as an unlabeled sample pool U; prepare a generative adversarial network as a synthetic sample generator; and prepare an image classification task model M to complete the image classification task.

[0010] Step S2: Train the generative adversarial network (GAN) from step S1 using the unlabeled sample pool U to obtain the sample generator G in the GAN; randomly select K samples from the unlabeled sample pool U as the first batch of data.

[0011] X1={x 11 x 12 ,..,x 1K}, train classification task model M;

[0012] Step S3: Randomly generate N random noises and input them into the synthetic sample generator G to obtain the i-th batch of generated samples.

[0013] Y i ={y i1 y i2 ,..,y iN}

[0014] Step S4: First, select a specific feature map F in the classification task model M. M Used for feature extraction and matching of samples; all samples in the unlabeled sample pool U are input into the classification task model M to obtain each sample x. ij Corresponding feature map F M (x ij For the sample set Y i Each sample y in ij Similarly, inputting it into the classification task model M yields the feature map F corresponding to each sample. M (y ij );

[0015] Step S5: Use Principal Component Analysis (PCA) to analyze the samples x obtained in Step S4. ij Corresponding feature map F M (x ijDimensionality reduction is performed; then, the K-means++ clustering algorithm is used to cluster the dimensionality-reduced data, obtaining K cluster centers; principal component analysis (PCA) is used to analyze the synthesized sample y. ij Feature map output F M (y ij Perform dimensionality reduction, calculate the distance between the dimensionality-reduced data and the K cluster centers, and obtain the nearest cluster center and its corresponding cluster a;

[0016] Step S6: Combine all samples x that are in cluster a ij With sample y ij Perform feature matching, i.e., calculate the feature difference score S. F (x ij y ij Then select the feature difference score S. F (x ij y ij The smallest sample a ij If the sample has already been matched by other synthesized samples, then the second-best sample is selected as sample a. ij ;

[0017] Step S7: Take the sample a obtained in step S6 ij With sample y ij Input the classification task model M for classification prediction; and use the class with the highest prediction confidence of the classification task model M for each of the two categories as its pseudo-labels, and calculate the gradient norm score T of the network parameters of several layers of the classification task model M. W And select sample a ij With sample y ij Gradient norm score T W The big ones j Add to the candidate sample pool;

[0018] Step S8: Repeat steps S6 to S7 until all N synthetic samples in the i-th batch of generated sample set have been iterated.

[0019] Step S9: Use a clustering algorithm to find the K cluster centers of the candidate sample pool and select the sample s that is closest to each cluster center. i Manually label the samples and add them to the labeled sample pool L. If sample s i If the sample is an existing sample in the unlabeled sample pool U, then remove the sample from the pool;

[0020] Step S10: Train the classification task model M using samples from the labeled sample pool L;

[0021] Step S11: Repeat steps S3 to S9 until the number of iterations reaches a preset value or the performance of the task model reaches the expected performance; finally, use the classification task model M to perform classification on the image classification dataset.

[0022] Furthermore, the generative adversarial network mentioned in step S1 is StyleGAN, and the classification task model M is ResNet-18.

[0023] Furthermore, the specific feature map F in step S4 M This is the output feature map of the Conv4_2 layer in the classification task model M.

[0024] Furthermore, in step S6, the feature difference score S between sample x and sample y is... F The formula for calculating (x, y) is as follows:

[0025]

[0026] The feature map size is C*H*W, where C corresponds to the number of channels in the feature map, H corresponds to the height of the feature map, and W corresponds to the width of the feature map; x ij y ij For the feature map components of each channel in the feature map of samples x and y, S is the pixel in the i-th row and j-th column. Fc (x, y) represents the feature difference scores for each channel of the feature map.

[0027] Furthermore, step S7 specifically includes:

[0028] Step S701: Set the classification task model M to gradient-free mode, and set the sample y ij Matching sample a ij Input the classification task model M, and obtain the classification task model M for sample y. ij Matching sample a ij The predicted class probability is used, and the class with the highest confidence is used as the pseudo-label of the sample.

[0029] Step S702: Adjust the task model to training mode, and then set the sample y ij Matching sample a ij Input the task model, obtain the gradients of its Conv5_2 layer and fully connected classification layer network parameters, and calculate the gradient norm score T. W and the gradient norm score T W The larger ones are added to the candidate sample pool as candidate samples; where the network parameter gradient norm score T W The calculations used the Conv5_2 layer and the final fully connected classification layer of the classification task model M as network parameters. The calculation method selected the L2 norm of each network parameter, and the specific calculation formula is as follows:

[0030]

[0031] Where W represents the gradient values ​​of all selected network layers, w i Let be the gradient values ​​of each parameter.

[0032] Furthermore, step S9 specifically includes:

[0033] Step S901: Input the samples in the candidate sample pool into the classification task model M to obtain the Conv4_2 input feature map, and use the principal component analysis method to reduce the dimensionality;

[0034] Step S902: Use the K-means++ clustering algorithm to cluster the dimensionality-reduced data, and set the number of cluster centers to K;

[0035] Step S903: Select K samples that are closest to each cluster center, which are the final samples selected for this round of active learning.

[0036] The diversity active learning method based on generative adversarial networks of the present invention has the following advantages:

[0037] This invention proposes a Generative-Based Diversity-Driven Active Learning (GDAL) strategy and applies it in practice to a railway train fault detection dataset, combined with a classification task model. First, this strategy introduces the high-performance Generative Adversarial Network StyleGAN as a synthetic sample generator, achieving high-resolution sample generation in complex scenarios. Second, for the original samples and GDAL-synthesized samples in the sample pool with differentiated feature information, the invention adopts an adaptive selection decision, utilizing the gradient norm score to adaptively select samples. This achieves diversified expansion of the sample pool and improves model robustness, while avoiding outlier specialization and sample invalidation problems that are prone to occur with purely synthetic sample-based methods. The adaptive selection strategy only uses the task model for decision calculation, avoiding the additional training cost of feature extractors such as autoencoders or discriminators. Finally, this strategy incorporates consideration of sample distribution diversity, reducing sample redundancy and further avoiding outlier specialization problems in the labeled sample pool. Experimental results show that the strategy in this paper achieves excellent decision performance on the railway train fault detection dataset, effectively contributing to reducing project labeling costs. Attached Figure Description

[0038] Figure 1 This is a schematic diagram of the process framework of the present invention.

[0039] Figure 2This is a schematic diagram illustrating the feature extraction and matching process using the feature map of the task model in step S4.

[0040] Figure 3 A diagram illustrating the brake hose, chain, and pin fault datasets for railway train fault detection.

[0041] Figure 4 A diagram illustrating the rolling bearing and axle end fault dataset for railway train fault detection.

[0042] Figure 5 This invention is used for performance comparison in a dataset of brake hose lifting chains and pin failures.

[0043] Figure 6 This is a performance comparison of the application of the present invention in a dataset of rolling bearing and shaft end faults. Detailed Implementation

[0044] To better understand the purpose, structure, and function of this invention, the following detailed description of a diversity active learning method based on generative adversarial networks is provided in conjunction with the accompanying drawings.

[0045] A Generative-Based Diversity-Driven Active Learning (GDAL) strategy based on generative adversarial networks, combined with Figure 1 The strategy flow diagram shown below illustrates the specific implementation steps as follows:

[0046] Step S1: Prepare a railway train fault detection image classification dataset, such as... Figure 3 and Figure 4 The datasets shown are for brake hose chain and pin faults and rolling bearing and shaft end faults, with 6000 and 8000 samples respectively, both 256*256 pixels in size. These will be used as the classification task dataset and the unlabeled sample pool U. A generative adversarial network, StyleGAN, will be prepared as the synthetic sample generator, and an image classification model, M = ResNet-18, will be prepared to complete the image classification task for railway train fault detection.

[0047] Step S2: Train a generative adversarial network using the classification dataset U to obtain the generator G. Randomly select K samples from the dataset, here K = 4% of the total data as the first batch of data X1 = {x l1 x 12 ,..,x 1K}, train classification task model M;

[0048] Step S3: Randomly generate N random noise samples, where N is set to 40% of the total data. Input these samples into generator G to obtain the generated sample set Y.i ={y i1 y i2 ,..,y iN}

[0049] Step S4: For the sample set Y i Each sample y in ij Inputting it into the task model ResNet-18 yields the selected feature map F of ResNet-18. M Perform sample matching. Select the difference scores S from the unlabeled sample pool U. F Minimum sample x ij ;

[0050] The feature extraction and matching method described in step S4 uses the Conv4_2 feature map output of ResNet-18 and introduces the K-means++ clustering algorithm and principal component analysis (PCA) method for feature matching, reducing the time consumed in the feature matching step; the feature difference score S between the two samples F The calculation formula is as follows:

[0051]

[0052] The feature map size is C*H*W, x ij y ij S represents the pixels corresponding to each channel in the feature maps of the two types of samples. Fc Calculate the feature difference scores for each channel. Step S4 specifically includes:

[0053] Step S401: First, input all samples in the unlabeled sample pool U into the task model ResNet-18, extract its Conv42 output feature map, and use the principal component analysis (PCA) method to reduce the data dimensionality.

[0054] Step S402: Cluster the dimensionality-reduced data using the K-means++ clustering algorithm;

[0055] Step S403: Input the synthetic samples generated by generator G into the task model ResNet-18, extract its Conv4_2 output feature map, and use principal component analysis (PCA) to reduce the data dimensionality.

[0056] Step S404: Perform K-means++ clustering on the dimensionality-reduced data to obtain the cluster category;

[0057] Step S405: Within the cluster's category, match the sample with the smallest difference score. If this sample has already been matched by other synthetic samples, select the next best sample.

[0058] Step S5: Pair task model M with x ijWith y ij The highest prediction confidence category is used as its pseudo-label, and the gradient norm score T of several layers of network parameters in the task model is calculated. W And select the one with the larger score. j Add to the candidate sample pool;

[0059] In step S5, the gradient norm score of the network parameters was calculated using the Conv5_2 layer and the final fully connected classification layer of the ResNet-18 task model as network parameters. The calculation method used was the L2 norm of each network parameter, and the specific calculation formula is as follows:

[0060]

[0061] Where W represents the gradient values ​​of all selected network layers, w i Let these be the gradient values ​​of each parameter. Step S5 specifically includes:

[0062] Step S501: Set the task model to gradient-free mode, input the synthetic sample and its matching sample into the task model, obtain the predicted class probability of the sample by the task model, and use the class with the highest confidence as the pseudo-label of the sample.

[0063] Step S502: Adjust the task model to training mode, then input the sample into the task model, obtain the gradient of the Conv5_2 layer and the fully connected classification layer network parameters and calculate the gradient norm, and add the sample with the larger gradient norm to the candidate sample pool.

[0064] Step S6: Use a clustering algorithm to find the K cluster centers of the candidate sample pool and select the sample s closest to each cluster center. i Manually label the samples and add them to the labeled sample pool L. If sample s i If the sample is an existing sample in the unlabeled sample pool, then remove the sample from the pool;

[0065] The clustering algorithm in step S6 employs the K-means++ clustering algorithm and uses image feature information for clustering. Similar to step S4, principal component analysis (PCA) is also used for data dimensionality reduction in this step. The feature maps used are the Conv4_2 output of ResNet-18. Step S6 specifically includes:

[0066] Step S601: Input the samples from the candidate sample pool into the task model to obtain the Conv4_2 input feature map, and use the principal component analysis (PCA) method to reduce the dimensionality;

[0067] Step S602: Use the K-means++ algorithm to cluster the dimensionality-reduced data, and set the number of cluster centers to K;

[0068] Step S603: Select the K samples closest to each cluster center, which are the final samples selected in this round of active learning;

[0069] Step S7: Train the task model M using samples from the labeled sample pool L;

[0070] Step S8: Repeat steps S3 to S7, stopping when the number of iterations has been pre-set or the task model performance reaches the expected result. Finally, the task model is used for classification of the railway train fault detection image classification dataset.

[0071] Regarding the training parameter settings for the ResNet task model, the number of iterations is set to 100, the optimizer is set to stochastic gradient descent (SGD), and the learning rate is l. r Set to 0.01, batch size to 512. Active learning iteration rounds to 10.

[0072] It is understood that the present invention has been described through some embodiments, and those skilled in the art will recognize that various changes or equivalent substitutions can be made to these features and embodiments without departing from the spirit and scope of the invention. Furthermore, under the teachings of the present invention, these features and embodiments can be modified to adapt to specific situations and materials without departing from the spirit and scope of the invention. Therefore, the present invention is not limited to the specific embodiments disclosed herein, and all embodiments falling within the scope of the claims of this application are within the protection scope of the present invention.

Claims

1. A diversity active learning method based on a generative adversarial network, characterized in that, The method comprises the following steps: Step S1: prepare an unlabeled image classification dataset and take it as an unlabeled sample pool U, prepare a generative adversarial network as a synthetic sample generator, and prepare an image classification task model M for completing an image classification task; Step S2: train the generative adversarial network in step S1 using the unlabeled sample pool U to obtain a synthetic sample generator G in the generative adversarial network; Randomly select K samples in the unlabeled sample pool U as the first batch of data X1 = {x 11 ,x 12 ,..,x 1K}, train a classification task model M; Step S3: randomly generate N random noises and input them into the synthetic sample generator G to obtain the i-th batch of generated sample set Y i = {y i1 , y i2 ,.., y iN} Step S4: First select a specific feature map F in the classification task model M M Feature extraction and matching for samples; input all samples in the unlabeled sample pool U into the classification task model M to obtain each x ij Corresponding feature map F M (x ij ); for each y i in the sample set Y ij , also input it into the classification task model M to obtain the corresponding feature map F M (y ij ) of each sample; x ij is the i-th row and j-th column pixel of each channel feature map component in the feature map of sample x; y ij is the i-th row and j-th column pixel of each channel feature map component in the feature map of sample y; Step S5: Use Principal Component Analysis (PCA) to analyze the x values ​​obtained in Step S4. ij Corresponding feature map F M (x ij Dimensionality reduction is performed; then, the K-means++ clustering algorithm is used to cluster the dimensionality-reduced data, obtaining K cluster centers; principal component analysis (PCA) is used to synthesize y ij Feature map output F M (y ij Perform dimensionality reduction, calculate the distance between the dimensionality-reduced data and the K cluster centers, and obtain the nearest cluster center and its corresponding cluster a; Step S6: Combine all x values ​​in cluster a ij With y ij Perform feature matching, i.e., calculate the feature difference score S. F (x ij ,y ij Then select the feature difference score S. F (x ij ,y ij The smallest sample a ij If the sample has already been matched by other synthesized samples, then the second-best sample is selected as sample a. ij ; Step S7: sample a obtained in step S6 is classified by the input classification task model M to obtain a classification result. ij with y ij inputting the classification task model M to make a classification prediction; And the highest prediction confidence class of the classification task model M for the two is respectively taken as its pseudo label, and the gradient norm score T of several layer network parameters of the classification task model M is calculated W And the sample a ij And y ij The larger gradient norm score T W Is selected j Into the candidate sample pool; Step S8: repeat the iteration of steps S6 to S7 until the N synthetic samples in the i-th batch of generated sample set are iterated completely; Step S9: Obtain K cluster centers of the candidate sample pool by using a clustering algorithm and select the sample s closest to each cluster center i Perform manual annotation and add the annotated sample pool L. If the sample s is an original sample in the unannotated sample pool U, remove the sample from the pool. i Perform manual annotation and add the annotated sample pool L. If the sample s is an original sample in the unannotated sample pool U, remove the sample from the pool. Step S10: train the classification task model M using the samples in the labeled sample pool L; Step S11: repeat the iteration of steps S3 to S9, and stop the iteration when the iteration number reaches a preset value or the performance of the task model reaches an expected performance; finally, the classification task model M is used for classification application in the image classification dataset. 2.The method of claim 1, wherein, The generative adversarial network in step S1 is StyleGAN, and the classification task model M is ResNet-18. 3.The diversity active learning method based on a generative adversarial network according to claim 2, characterized in that, The specific feature map F in step S4 M is the output feature map of the Conv4_2 layer in the classification task model M.

4. The generative adversarial network-based diversity active learning method according to claim 3, characterized in that, In step S6, the feature difference score S between the sample x and the sample y is calculated F The formula for (x, y) is as follows: wherein the feature map size is C*H*W, C corresponds to the channel number of the feature map, H corresponds to the height of the feature map, and W corresponds to the width size of the feature map; S Fc (x, y) is the feature difference score on each channel of the feature map. 5.The method of claim 1, wherein, Step S7 specifically comprises: Step S701: Set the classification task model M to a non-gradient mode, and input y ij and its matching sample a ij , input the classification task model M, and obtain the prediction class probability of the classification task model M on y ij and its matching sample a ij , and take the highest confidence class as the pseudo label of the sample; Step S702: adjust the task model to the training mode, and then obtain y ij The matching sample a ij Input the task model, obtain the gradient of the network parameters of the Conv5_2 layer and the fully connected classification layer, and calculate the gradient norm score T W , and the gradient norm score T W The larger one is added to the candidate sample pool; wherein the network parameter gradient norm score T W The calculation selects the Conv5_2 layer and the last fully connected classification layer of the classification task model M as the network parameters, and the L2 norm of each network parameter is selected as the calculation method. The specific calculation formula is as follows: where W is the selected all network layer gradient value, w i is where each parameter gradient value. 6.The method of claim 1, wherein, Step S9 specifically comprises: Step S901: input the samples in the candidate sample pool into the classification task model M to obtain Conv4_2 input feature maps, and perform dimension reduction by using a principal component analysis method; Step S902: cluster the dimension-reduced data by using a K-means++ clustering algorithm, and set the cluster centers to K; Step S903: select K samples closest to each cluster center as the final selected samples for the current round of active learning.

Images