A safety helmet detection method based on multi-stage personalized federated learning

Abstract

The application relates to a safety helmet detection method based on multi-stage personalized federated learning. The method divides the modeling process of safety helmet data owners in different scenes into four stages: model initialization, global model optimization, client model personalization and client model optimization. First, an initialization model is constructed through a public safety helmet detection data set. Second, a regular term describing the difference between the global safety helmet detection model and the client model is introduced in the common modeling process, thereby obtaining a global safety helmet detection model with high generalization. Third, the backbone network of the global safety helmet detection model is frozen and personalized migration is performed. Finally, the backbone network of the personalized safety helmet detection model is unfrozen and the model is optimized again. Compared with the prior art, the application solves or partially solves the problem that the existing method lacks sufficient learning and training data samples locally, and the model is difficult to train and has low precision.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of federated learning, and in particular to a safety helmet detection method based on multi-stage personalized federated learning. BACKGROUND

[0002] Safety helmets are important equipment for ensuring the safety of workers in production, and can effectively prevent and avoid accidental injuries to workers during production. With the continuous development and application of deep learning-based target detection algorithms, using intelligent recognition technology to replace manual supervision is a major trend. Currently, deep learning-based target detection models often require a large number of training samples to achieve good detection results. However, in the actual field, due to the possibility of safety helmet detection data leaking enterprise privacy, as well as national management measures and requirements for private data, safety helmet detection data usually exists in the form of an island, and it is difficult to collect a large amount of safety helmet detection data to train a model with high accuracy and high robustness. Federated learning, as a distributed machine learning technology that can protect user data privacy and security, can enable multiple parties to participate in modeling without sharing data, thereby improving the performance of each party's model. Although federated learning algorithms can simultaneously collaborate with multiple parties to establish multiple high-performance models, when the data owned by each party does not satisfy the independent and identically distributed characteristic, the performance of the traditional federated averaging algorithm will be greatly reduced, and the model may not converge at all.

[0003] On the one hand, existing safety helmet detection algorithms only collect data and train models for a certain scenario, and when there is a lack of a large number of data samples locally, the target detection model often fails to converge. On the other hand, existing personalized federated learning methods, such as data augmentation-based, client selection-based, and meta-learning-based methods, either risk leaking data privacy, or are computationally expensive, or require the selection of representative data sets first, so there is still a lot of room for improvement in the performance of personalized federated learning algorithms. Moreover, current personalized federated algorithms mostly target simple classification models, and for complex safety helmet detection models, how to efficiently collaborate with each party to train the model is also a major challenge.

[0004] The scenarios of safety helmet detection are usually very different, so it is necessary to design a federated learning algorithm with certain personalization for the participants to ensure the performance of each party's model, so that the safety helmet detection is accurate and effective, and the production and life safety are protected. SUMMARY

[0005] The purpose of the present application is to overcome the defects of the prior art and provide a safety helmet detection method based on multi-stage personalized federated learning. The final safety helmet detection model is obtained through four stages of model initialization, global model optimization, client model personalization and client model tuning for safety helmet detection. Since the method trains the model by combining public and private safety helmet detection data, and the private safety helmet detection data does not need to be uploaded to the server, the performance of the local safety helmet detection model can be improved without leaking private data.

[0006] The purpose of the present application can be achieved by the following technical solutions:

[0007] The present application provides a safety helmet detection method based on multi-stage personalized federated learning, applied to a server, comprising the following steps:

[0008] Obtain public safety helmet detection data, establish a global detection model, and perform initial training on the global detection model based on the public safety helmet detection data to initialize the model. The trained global detection model is distributed to each client.

[0009] Obtain the model learned by each client based on private safety helmet detection data, and update the global detection model. The updated global detection model is distributed to each client. Repeat this step multiple times to optimize the global model.

[0010] After the global model personalization and global model tuning of each client, the processed global model is used for safety helmet detection.

[0011] As a preferred technical solution, the process of obtaining the model learned by each client includes the following steps:

[0012] Each client uses cross-entropy to describe the difference between each client model and the global safety helmet detection model based on its own safety helmet detection data for training. After training, the client model weight of each client is uploaded to realize the acquisition of the model learned by each client.

[0013] As a preferred technical solution, the update of the global detection model parameters is realized by the following formula:

[0014]

[0015] loss(x i ,y i )+η(-∑f last (w)logf i (w))→f i (w)

[0016] In the formula, f(w) represents the updated global detection model parameter, f i (w) represents the client learning

[0017] i i The safety helmet detection result and the corresponding label data during client training are represented by f last (w) represents the last average aggregated global detection model parameter, η represents the weight adjustment parameter of the global model and the client model, : = represents the definition, n represents the number of clients, loss() represents the loss function, → represents parameter update, and w represents model weight.

[0018] As a preferred technical solution, the global model personalization process includes the following steps:

[0019] The backbone network part of the global model is frozen, and the corresponding safety helmet detection data of the client is used for training, and the step is repeated multiple times to realize global model personalization.

[0020] As a preferred technical solution, the loss function in the global model personalization process is:

[0021]

[0022] In the formula, loss user is the loss function value, loss(·) is the loss function, f i (w) is the model after personalization of each client, f(w) is the model parameter after optimization of the global model, represents the Frobenius norm, σ represents the personalization difference adjustment parameter, C s and C u respectively represent the feature covariance matrix of the public safety helmet detection data D s and the private safety helmet detection data D n , and d is the dimension of the input data.

[0023] As a preferred technical solution, the tuning process includes the following steps:

[0024] The backbone network part of the global model is unfrozen, and the global detection model of each client is trained using the corresponding safety helmet detection data of the client to realize tuning.

[0025] As a preferred technical solution, the public safety helmet detection data is obtained based on the public data set SHWD, and the private safety helmet detection data is made based on the LabelImg tool.

[0026] ​​As a preferred technical solution, the global detection model is constructed based on a YOLOv5 model.

[0027] As a preferred technical solution, the initial training is specifically:

[0028] An Aadm optimizer is adopted, and a cosine annealing method is used to update a learning rate, so that the initial training is realized.

[0029] As a preferred technical solution, the server obtains the model learned by each client based on the private safety cap detection data from each client through a homomorphic encryption algorithm.

[0030] Compared with the prior art, the present application has the following advantages:

[0031] (1) The final safety cap detection model is obtained through four stages of model initialization, global model optimization, client model individualization and client model tuning, and the safety cap detection model is used for safety cap detection. Compared with the existing safety cap detection algorithm, the safety cap detection model is trained by using the federated learning method, and the global model individualization processing and the global model tuning processing are performed after the training, so that the performance of the safety cap detection model in different scenes can be improved, and the robustness of the model can be enhanced without sharing the private data.

[0032] (2) In the global model individualization processing and the global model tuning processing, the backbone network part of the global model is first frozen and individualized migrated, and after repeating the foregoing steps for multiple times, the backbone network part of the global model is unfrozen for model tuning, so that the individualization training effect is effectively improved.

[0033] (3) The multi-stage individualized federated learning algorithm provided by the present application can effectively reduce the influence of the safety cap samples in different scenes which are non-independent and identically distributed data, and compared with the existing individualized federated learning algorithm which can only improve the performance of a simple classification model, the present method is more conducive to the convergence of a complex target detection model. BRIEF DESCRIPTION OF DRAWINGS

[0034] Figure 1 It is a flowchart of the safety cap detection method based on the multi-stage individualized federated learning in embodiment 1.

[0035] Figure 2 It is a schematic diagram of the technical framework of the present application.

[0036] Figure 3 It is a schematic diagram of the implementation process in embodiment 1.

[0037] Figure 4 It is a schematic diagram of the safety cap detection data sample.

[0038] Figure 5 A safety helmet detection accuracy comparison chart. DETAILED DESCRIPTION

[0039] The technical solutions in the embodiments of the present application will be clearly and completely described below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are part of the embodiments of the present application, rather than all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by a person of ordinary skill in the art without creative labor should fall within the protection scope of the present application.

[0040] Embodiment 1

[0041] As Figure 1 described, in order to solve the problems that the existing safety helmet detection model and method lack sufficient learning training data samples locally, the model is difficult to train and the precision is low, and the safety helmet data samples in different scenarios are not independent and identically distributed, which leads to the problem that the safety helmet detection model is not easy to converge in the federated learning framework, the embodiment provides a safety helmet detection method based on multi-stage personalized federated learning.

[0042] As Figure 2 described, the schematic diagram of the technical framework of the embodiment is shown. First, define the safety helmet data owners in different scenarios as different participants, and define the participant set as U={U1,…,U n}, the safety helmet data set in different scenarios as D={D1,…,D n}, and any participant U i will not share the owned data D i with other participants. Define the model set of different participants as the client model set A u ={A1,…,A n}. The server U s owns the public safety helmet detection data D s . The goal of the safety helmet detection method based on multi-stage personalized federated learning proposed in the present application is to make the safety helmet detection model precision of any client close to the model trained using the whole data. It can be described by the following formula (1) formally, where A all represents the model trained using all the data D, mAP represents the model precision, and Δ is a small positive real number.

[0043]

[0044] The safety helmet detection method includes the following four stages:

[0045] Step S1, model initialization: an initialization model is constructed through a public safety helmet detection data set. The server U sUsing public security cap detection data D s An initialization model is trained, then distributed to each participant, and finally the weights of the initialization model are used as the pre-training weights of the client models.

[0046] Step S2, global model optimization: a regularization term is introduced to describe the difference between the global security cap detection model and the client model in the common modeling process, so as to obtain a global security cap detection model with higher generalization. All participants use the security cap detection data they have, upload their client model weights after 5 rounds of training, and then the server side obtains a global security cap detection model through average aggregation and distributes it to each participant. To prevent the client model from deviating during the global model optimization stage, the cross-entropy is used to describe the difference between the client model and the global security cap detection model, and it is introduced into the client model training process. Therefore, the model parameter generation method in the global model optimization stage can be represented as:

[0047]

[0048] Where f(w) represents the global security cap detection model parameters, f i (w) represents each client model. f i (w) can be calculated by formula (3), where x i ,y i represents the security cap detection result and the corresponding label data during client training, f last (w) represents the global security cap detection model parameters of the last average aggregation, and η represents the weight adjustment parameter of the global model and the client model.

[0049] loss(x i ,y i )+η(-∑f last (w)logfi i (w))→f i (w) (3)

[0050] Step S3, client model personalization: freeze the backbone network of the global security cap detection model and perform personalized migration. Since the backbone network of the security cap detection model is used to extract the features required for detection, it has universality, so the backbone network of the client model is first frozen to speed up the personalization of the client model. The multi-scale fusion network and detection head network of the target detection network are high-level features associated with the detection task and the participant data, and in the personalization stage, these parameters are mainly updated. Finally, in order to enable the client model to learn the difference between the client and the server, the client model loss function is defined as follows:

[0051]

[0052] where C s , C u denote the public data D s owned by the clients and the feature covariance matrix of the client data D n , respectively. denotes the Frobenius norm, and σ denotes the personalization difference adjustment parameter.

[0053] Step S4, client model tuning: unfreeze the backbone network of the personalized safety hat detection model, and retrain the model. In this stage, the backbone model parameters are no longer frozen, and the personalized client model is trained again using the client data.

[0054] In summary, the safety hat detection method based on multi-stage personalized federated learning has the following steps:

[0055] Method input: safety hat data set D in different scenarios, public data set d s , η, σ.

[0056] Method output: personalized client model.

[0057] Step 1, use the public data set D s to build an initial model f(w) and distribute f(w) to all participants U.

[0058] Step 2, the global model optimization stage starts, and each client model is trained using D and formula (3).

[0059] Step 3, upload the client model through homomorphic encryption, update the global model through formula (2), and distribute it to all participants U.

[0060] Step 4, repeat Step 2-3 until the global model optimization stage ends.

[0061] Step 5, the client model personalization stage starts, and the backbone network part of the client model is frozen, and the client model is personalized using formula (4).

[0062] Step 6, repeat Step 5 until the client model personalization stage ends.

[0063] Step 7, the client tuning stage starts, unfreezes the backbone network of the client model, and tunes the personalized client model using local data.

[0064] A typical implementation of the above method is provided below, mainly including two aspects of data preparation and programming implementation of the safety helmet detection method of multi-stage personalized federated learning. Specifically as follows:

[0065] (1) Data preparation

[0066] As shown in the following Figure 2 The present application uses the public dataset SHWD as the server-side data, and the pictures in the dataset are mostly from the network, a total of 7581 images, of which 9044 samples wear safety helmets and 111514 samples do not wear safety helmets. Then each participant collects private safety helmet detection data and uses the LabelImg tool to label the collected data to form multiple private safety helmet detection datasets.

[0067] (2) Programming implementation of the safety helmet detection method of multi-stage personalized federated learning

[0068] First, the YOLOv5 model is selected as the base model, and the code is implemented using the PyTorch framework. In the initialization stage, the SHWD dataset is selected to train the YOLOv5 model, and the training settings are as follows: use the Aadm optimizer, define the initial learning rate as 0.00001, the batch size as 16, and use the cosine annealing method to update the learning rate, and train a total of 300 rounds.

[0069] After completing the initialization model training, the model is distributed to each participant and used as a client model. Subsequently, each participant uses private safety helmet detection data and utilizes the loss function of the global model optimization stage to train the model, and the training settings are as follows: use the SGD optimizer, define the initial learning rate as 0.001, the weight decay as 5e-4, the batch size as 32, and train a total of 100 rounds. And every 5 rounds of training, each participant pauses the training, then encrypts the weight value of the client model using the homomorphic encryption algorithm and sends it to the server side. The server side decrypts the weight value of each client model, then aggregates the weight value of all client models. Finally, the aggregated weight value is distributed to each participant, and the weight parameter is used as the weight of the client model for further training.

[0070] After completing the global model optimization, the backbone network of the client model is frozen, and then the client model is personalized trained using private safety helmet detection data and the loss function of the personalized stage, and the training settings are as follows: use the SGD optimizer, define the initial learning rate as 0.001, the weight decay as 5e-4, the batch size as 32, and train a total of 100 rounds.

[0071] In the final tuning stage, the backbone network of the unfreezing client model is used to continue training the safety helmet detection model after the completion of the client model personalization using private safety helmet detection data, and the training settings are as follows: using the SGD optimizer, defining the initial learning rate as 0.001, the weight decay as 5e-4, the batch size as 16, and a total of 300 rounds of training.

[0072] As shown in Figure 4 , the application demonstration of the present application is carried out using safety helmet detection data from four scenarios of a steel plant (sample 1), an automobile plant (sample 2), a chemical plant (sample 3), and a power plant (sample 4), and η = 0.001, σ = 0.005 are set. As shown in Figure 5 , the safety helmet detection method of the multi-stage personalized federated learning proposed by the present application can effectively improve the accuracy of the safety helmet detection model of each participant.

[0073] Compared with existing safety helmet detection algorithms, the present application can improve the performance of safety helmet detection models in different scenarios without sharing private data. At the same time, the multi-stage personalized federated learning algorithm proposed can effectively reduce the impact of safety helmet samples in different scenarios as non-identically distributed data, and compared with the existing personalized federated learning algorithm which can only improve the performance of simple classification models, the present method is more conducive to the convergence of complex target detection models.

[0074] Embodiment 2

[0075] The present embodiment provides an electronic device, comprising one or more processors and a memory, the memory having one or more programs stored therein, the one or more programs comprising instructions for performing the safety helmet detection method based on multi-stage personalized federated learning as described in Embodiment 1.

[0076] Embodiment 3

[0077] The present embodiment provides a computer-readable storage medium comprising one or more programs for execution by one or more processors of an electronic device, the one or more programs comprising instructions for performing the safety helmet detection method based on multi-stage personalized federated learning as described in Embodiment 1.

[0078] The above is merely a specific implementation of the present application, but the scope of protection of the present application is not limited thereto, and any person skilled in the art can easily think of various equivalent modifications or replacements within the technical scope disclosed by the present application, and these modifications or replacements should be covered within the scope of protection of the present application. Therefore, the scope of protection of the present application should be subject to the scope of protection of the claims.

Claims

1. A hard hat detection method based on multi-stage personalized federated learning, characterized in that, The application is applied to a server and includes the following steps: Obtaining public safety helmet detection data, establishing a global detection model, performing initial training on the global detection model based on the public safety helmet detection data, completing initialization of the model, and distributing the trained global detection model to each client; Obtaining the model learned by each client based on private safety helmet detection data, updating the global detection model, distributing the updated global detection model to each client, repeating the step multiple times to optimize the global model; After individualization and optimization of the global model, each client uses the processed global model to detect safety helmets. The global model individualization process includes the following steps: Freezing the backbone network part of the global model, training it using the safety helmet detection data corresponding to the client, and repeating the step multiple times to achieve global model individualization. The optimization process includes the following steps: Unfreezing the backbone network part of the global model, training it using the safety helmet detection data corresponding to the client, and achieving optimization.

2. The safety helmet detection method based on multi-stage personalized federated learning according to claim 1, characterized in that, The process of obtaining the learned model of each client includes the following steps: Each client uses cross-entropy to describe the difference between the client model and the global safety helmet detection model based on its own safety helmet detection data, trains the model, uploads the client model weight after training, and obtains the learned model of each client.

3. The method of claim 1, wherein the method is based on a multi-stage personalized federated learning. The update of the global detection model parameters is realized by the following formula: In the formula, represents the updated global detection model parameters, represents the model learned from each client, represents the safety cap detection result during client training and the corresponding label data, represents the last average aggregated global detection model parameters, represents the weight adjustment parameters of the global model and the client model, represents the definition, represents the number of clients, represents the loss function, represents the parameter update, represents the model weight.

4. The safety hat detection method based on multi-stage personalized federated learning according to claim 1, characterized in that, In the global model individualization process, the loss function is: wherein, is a loss function value, is a loss function, is a model personalized for each client, is a model parameter optimized by the global model, denotes the Frobenius norm, denotes a personalized difference adjustment parameter, , denote the feature covariance matrix of public safety hat detection data and private safety hat detection data respectively, is the dimension of the input data.

5. The method of claim 1, wherein the method is a method of safety hat detection based on multi-stage personalized federated learning. The public safety helmet detection data is obtained based on the public data set SHWD, and the private safety helmet detection data is obtained based on the LabelImg tool.

6. The method of claim 1, wherein the method is a method of hard hat detection based on multi-stage personalized federated learning. The global detection model is constructed based on the YOLOv5 model.

7. The method of claim 1, wherein the method is a method of hard hat detection based on multi-stage personalized federated learning. The initial training is as follows: Using the Aadm optimizer and the cosine annealing method to update the learning rate to realize the initial training of the global detection model.

8. The safety hat detection method based on multi-stage personalized federated learning according to claim 1, characterized in that, The server obtains the model learned by each client based on the private safety helmet detection data through the homomorphic encryption algorithm.

Images