Steel surface defect classification method and system based on adversarial personalized federated learning

Abstract

The present invention provides a steel surface defect classification method and system based on adversarial personalized federated learning, which belongs to the technical field of steel surface defect classification, and includes the following steps: participating users collect and annotate local steel surface defect samples to construct a local data set; a central server distributes a global encoder; participating users update the initial local encoder, initial local classifier and initial local discriminator based on a dynamic consensus construction strategy; participating users upload the first local encoder and identification loss; the central server obtains the next round of global encoders based on a consensus-aware aggregation mechanism; participating users update the first local encoder, the first local classifier and fusion weights based on an adaptive feature fusion module; and repeat the above-mentioned federated communication process until the preset maximum communication round. The present invention solves the problem of poor model performance caused by data dispersion and data heterogeneity in the industrial steel surface defect classification scenario.

Description

Technical Field

[0001] The present invention relates to the field of steel surface defect classification, and in particular to a steel surface defect classification method and system based on adversarial personalized federated learning. Background Art

[0002] Surface defects are a major concern for the quality and safety of manufactured materials. Traditional surface defect classification methods mainly rely on manual inspection, which is not only time-consuming but also error-prone, affecting production efficiency and the accuracy of quality control. With the development of artificial intelligence technology, especially the application of deep learning technology, new solutions are provided for automated surface defect detection. Efficient and accurate defect detection can be achieved by using deep learning architectures based on convolutional neural networks, which are able to autonomously learn abstract representations of data, thus eliminating the need for manual feature engineering.

[0003] However, deep learning-based methods face a major problem, which is that a large amount of data is required to train the model. It is difficult for a single production entity to collect enough training data, and aggregating data from multiple entities raises the issue of data privacy. Federated learning, as an emerging distributed machine learning paradigm, effectively solves the problems of data scarcity and privacy protection by training models locally on the client and sharing only model parameters rather than the data itself. However, the implementation of federated learning is still challenging when faced with data heterogeneity. Due to differences in production conditions among different industrial entities, the data distribution of each client is inconsistent, which may affect the performance of the global model.

[0004] Therefore, there is an urgent need for a distributed training method for the data heterogeneity problem in steel surface defect classification scenarios, which can effectively overcome the problem of global knowledge forgetting in local training caused by data heterogeneity, enhance the global knowledge learning effect of local users in heterogeneous environments, and train a personalized model with superior performance for each user while ensuring privacy security. Summary of the invention

[0005] The purpose of the present invention is to provide a steel surface defect classification method and system based on adversarial personalized federated learning, which promotes multiple clients to collaboratively train personalized models for each client's specific data distribution through adversarial training and using a global model as a bridge, while protecting privacy and effectively improving the adverse effects of data heterogeneity.

[0006] To achieve the above object, the present invention provides a steel surface defect classification method based on adversarial personalized federated learning, comprising the following steps:

[0007] S1, participating users collect and annotate local steel surface defect samples to build a local dataset;

[0008] S2, the central server distributes the global encoder;

[0009] S3, the participating users update the initial local encoder, the initial local classifier and the initial local discriminator based on the dynamic consensus building strategy to obtain the first local encoder, the first local classifier and the first local discriminator;

[0010] S4, the participating user uploads the first local encoder and the identification loss to the central server;

[0011] S5, the central server calculates the next round of global encoder based on the consensus-aware aggregation mechanism;

[0012] S6. The participating user updates the first local encoder, the first local classifier and the fusion weight based on the adaptive feature fusion module;

[0013] S7. Repeat steps S2 to S6 until the maximum communication round is reached.

[0014] Preferably, in step S3, the participating users update the initial local encoder, the initial local classifier and the initial local discriminator based on the dynamic consensus building strategy to obtain the first local encoder, the first local classifier and the first local discriminator. The specific operations are as follows:

[0015] S31. For defective sample x, participating user k inputs it into the global encoder respectively. With the initial local encoder Get the global feature corresponding to the defect sample x With local features t is the communication round;

[0016] S32, local features Input to the initial local classifier According to the corresponding specific defect category label, the classification loss is calculated by the cross entropy loss function The classification loss The specific calculation formula is as follows:

[0017]

[0018] in, is the local data set of participating user k; l is the sample number; is the defect sample image corresponding to the sample number l in the local dataset; is the corresponding category label; (x, y) represents the image and label pair in the local dataset; y is the label; N is the number of defect categories; 1 is the indicator function; is the initial local classifier The output probability distribution vector The probability value corresponding to defect category c;

[0019] S33, local features With global features Input to the initial local discriminator According to its corresponding local features Or global features The assigned domain label is used to calculate the identification loss using the cross entropy loss function The identification loss The calculation formula is as follows:

[0020]

[0021] in, It is a set of features and domain labels; local features and global features The domain labels are 0 and 1 respectively; is the initial local discriminator The output probability distribution vector The probability value corresponding to the defect category c in ; (f, y) is an element in the set consisting of features and domain labels;

[0022] S34. Minimize classification loss To update the initial local encoder With the initial local classifier Minimize the discrimination loss To update the initial local discriminator While maximizing the discrimination loss To update the initial local encoder as follows:

[0023]

[0024]

[0025]

[0026] Among them, α is the learning rate; λ is the discrimination loss trade-off coefficient; is a first local encoder; is the first local classifier; is the first local discriminator.

[0027] The dynamic consensus building strategy is based on adversarial training technology. While training the classification task, it shortens the distance between the global features and local features of the same sample, using the global encoder as a bridge to achieve distribution convergence between heterogeneous users, alleviate the problem of global knowledge forgetting during training under heterogeneous data conditions, enhance the global knowledge learning effect of local users, and improve model performance.

[0028] Preferably, in step S5, the central server calculates the next round of global encoders based on the consensus-aware aggregation mechanism, and the specific operations are as follows:

[0029]

[0030] Among them, K is the total number of participating users; The global encoder of round t+1 calculated by the server.

[0031] The consensus-aware aggregation mechanism evaluates the effect of user global knowledge learning by the size of identification loss, and assigns higher weights to user models with higher global knowledge learning effects, so as to generate a global model with higher generalization performance. It complements the dynamic consensus building strategy, further improves the effect and efficiency of consensus building, and thus improves model performance.

[0032] Preferably, in step S6, the participating user updates the first local encoder, the first local classifier and the fusion weight based on the adaptive feature fusion module, and the specific operations are as follows:

[0033] S61. For defective sample x, participating user k inputs it into the global encoder respectively. With the first local encoder Get the global feature corresponding to the defect sample x With local features

[0034] S62. Using trainable scalar fusion weights For global features With local features Perform feature fusion:

[0035]

[0036] Then the fusion features corresponding to all defect samples x Input first local classifier According to the corresponding specific defect category label, the classification loss is calculated by the cross entropy loss function. The classification loss The specific calculation formula is as follows:

[0037]

[0038] in, is the first local classifier The output probability distribution vector The probability value corresponding to defect category c;

[0039] S63. Minimize classification loss Update the first local encoder First local classifier and fusion weight as follows:

[0040]

[0041]

[0042]

[0043] in, is a second local encoder; is the second local classifier; is the updated fusion weight.

[0044] The adaptive feature fusion module adaptively learns the best personalized fusion weight for each user by gradually optimizing the weight when fusing the global features and local features of each user, thereby improving the utilization efficiency of global knowledge and thus improving model performance.

[0045] The present invention also provides a system for implementing the steel surface defect classification method based on adversarial personalized federated learning as described above, comprising a central server and a number of participating users with heterogeneous data distribution, wherein the local model of the participating users comprises an initial local encoder, an initial local classifier and an initial local discriminator.

[0046] Therefore, the present invention adopts the above-mentioned steel surface defect classification method and system based on adversarial personalized federated learning, and the technical effects are as follows:

[0047] (1) The federated learning framework is used to jointly train the model using the data of multiple participants while ensuring data privacy, which greatly improves the accuracy of the model;

[0048] (2) The dynamic consensus building strategy adopts adversarial training technology and uses the global model as a bridge to achieve the convergence of feature distributions between different users, alleviating the problem of data heterogeneity. On this basis, the consensus-aware aggregation strategy uses the discrimination loss values ​​of different users as aggregation weights to enhance the generalization ability of the global model and further improve the learning effect. Finally, the adaptive feature fusion module learns personalized global and local feature fusion modes for different users, further improving model performance. BRIEF DESCRIPTION OF THE DRAWINGS

[0049] Figure 1 A flow chart of a steel surface defect classification method and system based on adversarial personalized federated learning according to the present invention;

[0050] Figure 2 This is a diagram of experimental results of Example 1 of the present invention under two types of random heterogeneous settings;

[0051] Figure 3 This is a graph showing the experimental results of Example 1 of the present invention under a Dirichlet heterogeneous setting;

[0052] Figure 4 This is an ablation experiment diagram of different submodules in embodiment 1 of the present invention under two types of random heterogeneous settings;

[0053] Figure 5 This is an ablation experiment diagram of different submodules in Dirichlet heterogeneous setting in Example 1 of the present invention. DETAILED DESCRIPTION

[0054] The technical solution of the present invention is further described below through the accompanying drawings and embodiments.

[0055] Unless otherwise defined, technical or scientific terms used in the present invention shall have the common meanings understood by one having ordinary skills in the field to which the present invention belongs.

[0056] Embodiment 1

[0057] like Figure 1 As shown, the present invention provides a steel surface defect classification method based on adversarial personalized federated learning, comprising the following steps:

[0058] S1, participating users collect and annotate local steel surface defect samples to build a local dataset;

[0059] This embodiment uses a steel surface defect classification method based on adversarial personalized federated learning, and uses a personalized federated learning framework to jointly train a personalized model suitable for local data characteristics for each manufacturer based on the steel surface defect classification data of different manufacturers. The NEU steel surface defect dataset used is a public dataset, containing 1,800 samples of six defect categories: rolling scale, plaque, crack, concave surface, inclusions, and scratches.

[0060] In this embodiment, the number of categories or category marginal distribution owned by each manufacturer is different, and two data heterogeneous distributions are used: 1. Each manufacturer has two random defect samples; 2. The data is divided for each manufacturer according to the Dirichlet distribution, and the Dirichlet distribution parameter is set to 0.1. The number of local samples of each manufacturer is set to 5, 10, and 20 to verify the effectiveness of this method under different sample numbers. The local model of each manufacturer includes an encoder, a classifier, and a discriminator.

[0061] S2, the central server distributes the global encoder;

[0062] S3, the participating users update the initial local encoder, the initial local classifier and the initial local discriminator based on the dynamic consensus building strategy to obtain the first local encoder, the first local classifier and the first local discriminator;

[0063] S31. For defective sample x, participating user k inputs it into the global encoder respectively. With the initial local encoder Get the global feature corresponding to the defect sample x With local features t is the communication round;

[0064] S32, local features Input to the initial local classifier According to the corresponding specific defect category label, the classification loss is calculated by the cross entropy loss function The classification loss The specific calculation formula is as follows:

[0065]

[0066] in, is the local data set of participating user k; l is the sample number; is the defect sample image corresponding to the sample number l in the local dataset; is the corresponding category label; (x, y) represents the image and label pair in the local dataset; y is the label; N is the number of defect categories; 1 is the indicator function; is the initial local classifier The output probability distribution vector The probability value corresponding to defect category c;

[0067] S33, local features With global features Input to the initial local discriminator According to its corresponding local features Or global features The assigned domain label is used to calculate the identification loss using the cross entropy loss function The identification loss The calculation formula is as follows:

[0068]

[0069] in, It is a set of features and domain labels; local features and global features The domain labels are 0 and 1 respectively; is the initial local discriminator The output probability distribution vector The probability value corresponding to the defect category c in ; (f, y) is an element in the set consisting of features and domain labels;

[0070] S34. Minimize classification loss To update the initial local encoder With the initial local classifier Minimize the discrimination loss To update the initial local discriminator While maximizing the discrimination loss To update the initial local encoder as follows:

[0071]

[0072]

[0073]

[0074] Among them, α is the learning rate; λ is the discrimination loss trade-off coefficient; is a first local encoder; is the first local classifier; is the first local discriminator.

[0075] S4, the participating user uploads the first local encoder and the identification loss to the central server;

[0076] S5. The central server calculates the next round of global encoders based on the consensus-aware aggregation mechanism. The specific operations are as follows:

[0077]

[0078] Among them, K is the total number of participating users; The global encoder of round t+1 calculated by the server.

[0079] S6. The participating user updates the first local encoder, the first local classifier and the fusion weight based on the adaptive feature fusion module;

[0080] S61. For defective sample x, participating user k inputs it into the global encoder respectively. With the first local encoder Get the global feature corresponding to the defect sample x With local features

[0081] S62. Using trainable scalar fusion weights For global features With local features Perform feature fusion:

[0082]

[0083] Then the fusion features corresponding to all defect samples x Input first local classifier According to the corresponding specific defect category label, the classification loss is calculated by the cross entropy loss function. The classification loss The specific calculation formula is as follows:

[0084]

[0085] in, is the first local classifier The output probability distribution vector The probability value corresponding to defect category c;

[0086] S63. Minimize classification loss Update the first local encoder First local classifier and fusion weight as follows:

[0087]

[0088]

[0089]

[0090] in, is a second local encoder; is the second local classifier; is the updated fusion weight.

[0091] S7. Repeat steps S2 to S6 until the maximum communication round is reached.

[0092] This example uses the following methods for comparison:

[0093] (1) FedAvg: Federated Averaging Method, which is used as the baseline method for federated learning methods.

[0094] (2) FedProx: Adds proximal terms to the local loss function of FedAvg to alleviate the impact of data heterogeneity.

[0095] (3) FedRep: A personalized federated learning method that uses parameter decoupling to train models.

[0096] (4) Ditto: A personalized federated learning method that uses a multi-task learning approach for model training.

[0097] (5) FedALA: A personalized federated learning method that adaptively fuses the downloaded global model and local model to the local target of each client and initializes the local model before each iterative training.

[0098] The experimental results of defect classification are shown in Figure 2 and Figure 3 As shown, it can be seen that under two different data heterogeneity settings and three different training sample number settings, the present invention achieves higher classification accuracy compared with other methods.

[0099] The ablation experimental results of the method proposed in the present invention are as follows: Figure 4 and Figure 5 As shown in the figure, the three submodules included in this method are: dynamic consensus building strategy, consensus-aware aggregation strategy, and adaptive feature fusion module, which are represented by A, B, and C respectively. It can be seen that the gradual enhancement of performance between different variants and the complete method verifies the effectiveness of each module.

[0100] Embodiment 2

[0101] The present invention also provides a steel surface defect classification system based on adversarial personalized federated learning, including a central server and a number of participating users with heterogeneous data distribution, wherein the local model of the participating users includes a local encoder, a local classifier and a local discriminator.

[0102] Therefore, the present invention adopts the above-mentioned steel surface defect classification method and system based on adversarial personalized federated learning, and promotes multiple clients to collaboratively train personalized models for each client's specific data distribution through adversarial training and using the global model as a bridge, while protecting privacy and effectively improving the adverse effects of data heterogeneity.

[0103] Finally, it should be noted that the above embodiments are only used to illustrate the technical solution of the present invention rather than to limit it. Although the present invention has been described in detail with reference to the preferred embodiments, those skilled in the art should understand that they can still modify or replace the technical solution of the present invention with equivalents, and these modifications or equivalent replacements cannot cause the modified technical solution to deviate from the spirit and scope of the technical solution of the present invention.

Claims

1. A steel surface defect classification method based on adversarial personalized federated learning, characterized in that: The following steps are involved: S1, participating users collect and annotate local steel surface defect samples to build a local dataset; S2, the central server distributes the global encoder; S3, the participating users update the initial local encoder, the initial local classifier and the initial local discriminator based on the dynamic consensus building strategy to obtain the first local encoder, the first local classifier and the first local discriminator; S4, the participating user uploads the first local encoder and the identification loss to the central server; S5, the central server calculates the next round of global encoder based on the consensus-aware aggregation mechanism; S6. The participating user updates the first local encoder, the first local classifier and the fusion weight based on the adaptive feature fusion module; S7, repeat steps S2 to S6 until the maximum communication round is preset; In step S3, the participating users update the initial local encoder, the initial local classifier and the initial local discriminator based on the dynamic consensus building strategy to obtain the first local encoder, the first local classifier and the first local discriminator. The specific operations are as follows: S31. For defective sample x, participating user k inputs it into the global encoder respectively. With the initial local encoder Get the global feature corresponding to the defect sample x With local features t is the communication round; S32, local features Input to the initial local classifier According to the corresponding specific defect category label, the classification loss is calculated by the cross entropy loss function The classification loss The specific calculation formula is as follows: in, is the local data set of participating user k; l is the sample number; is the defect sample image corresponding to the sample number l in the local dataset; is the corresponding category label; (x, y) represents the image and label pair in the local dataset; y is the label; N is the number of defect categories; 1 [y＝c] is the indicator function; is the initial local classifier The output probability distribution vector The probability value corresponding to defect category c; S33, local features With global features Input to the initial local discriminator According to its corresponding local features Or global features The assigned domain label is used to calculate the identification loss using the cross entropy loss function The identification loss The calculation formula is as follows: in, It is a set of features and domain labels; local features and global features The domain labels are 0 and 1 respectively; is the initial local discriminator The output probability distribution vector The probability value corresponding to the defect category c in ; (f, y) is an element in the set consisting of features and domain labels; S34. Minimize classification loss To update the initial local encoder With the initial local classifier Minimize the discrimination loss To update the initial local discriminator While maximizing the discrimination loss To update the initial local encoder as follows: Among them, α is the learning rate; λ is the discrimination loss trade-off coefficient; is a first local encoder; is the first local classifier; is the first local discriminator.

2. According to claim 1, a steel surface defect classification method based on adversarial personalized federated learning is characterized in that: In step S5, the central server calculates the next round of global encoders based on the consensus-aware aggregation mechanism. The specific operations are as follows: Among them, K is the total number of participating users; The global encoder of round t+1 calculated by the server.

3. The steel surface defect classification method based on adversarial personalized federated learning according to claim 2 is characterized in that: In step S6, the participating user updates the first local encoder, the first local classifier and the fusion weight based on the adaptive feature fusion module. The specific operations are as follows: S61. For defective sample x, participating user k inputs it into the global encoder respectively. With the first local encoder Get the global feature corresponding to the defect sample x With local features S62. Using trainable scalar fusion weights For global features With local features Perform feature fusion: Then the fusion features corresponding to all defect samples x Input first local classifier According to the corresponding specific defect category label, the classification loss is calculated by the cross entropy loss function. The classification loss The specific calculation formula is as follows: in, is the first local classifier The output probability distribution vector The probability value corresponding to defect category c; S63. Minimize classification loss Update the first local encoder First local classifier and fusion weight as follows: in, is a second local encoder; is the second local classifier; is the updated fusion weight.

4. A system for implementing the steel surface defect classification method based on adversarial personalized federated learning as described in any one of claims 1 to 3, characterized in that: The method comprises a central server and a number of participating users with heterogeneous data distribution, wherein the local models of the participating users comprise an initial local encoder, an initial local classifier and an initial local discriminator.

Images