Defect recognition model lightweight method and system based on reinforcement learning

Abstract

The application discloses a kind of based on reinforcement learning's defect identification model lightening method and system, the method includes according to defect identification model construction sample data set;The i th sample is used to the i th time update policy network and value network, obtain the i th time updated policy network and value network;According to the i th time updated policy network and total target pruning proportion, defect identification model is pruned, and the identification accuracy of pruned defect identification model on verification set is calculated;After completing the training of policy network and value network, according to the trained policy network and total target pruning proportion, defect identification model is pruned, and the lightened defect identification model is obtained.The application reduces the parameter quantity of defect identification model on the basis of ensuring defect identification accuracy, realizes the balance of the identification performance and lightening of defect identification model.

Description

TECHNICAL FIELD

[0001] The application belongs to the field of artificial intelligence and computer technology, and particularly relates to a defect recognition model lightweight method and system based on reinforcement learning. BACKGROUND

[0002] As an efficient, flexible and safe aerial sensing platform, unmanned aerial vehicles (UAVs) have been widely used in intelligent inspection tasks of power transmission lines, tower bodies, insulators, fittings and conductors. By carrying high-resolution cameras, UAVs can quickly complete image acquisition in complex, dangerous and difficult-to-reach environments for humans, and complete defect recognition and intelligent alarm through background algorithms, significantly improving inspection efficiency and operation safety, and have become a key technical means in the modern intelligent operation and maintenance system of power grids.

[0003] In the task of intelligent inspection of power, image defect recognition models based on deep learning (referred to as deep learning models) play a crucial role. Compared with traditional rule-based image processing methods, deep learning models have powerful feature extraction and pattern recognition capabilities, and can accurately identify various common defects such as damage, foreign object attachment and bird nests. However, in order to obtain higher accuracy, many deep learning models tend to be complex in structure design, resulting in large parameter size, high computational cost and high inference delay, which limits the actual deployment and operation efficiency of deep learning models on the UAV side.

[0004] In the intelligent identification task of power defects, the UAV side, the edge side and the cloud side often work together. For some simple tasks, such as obvious features and easily identifiable target categories, a lightweight recognition model can be deployed on the UAV side to identify and feedback in real time. When it is difficult to identify, the preliminary identification on the UAV side is not successful or the credibility is insufficient, it is further transmitted to the edge side and the cloud side for further identification by using high-precision and more complex models. Due to the limited computing resources, storage space and battery life of the UAV side, deep learning models must balance the dual requirements of performance and lightweight to perform real-time inference on the UAV side.

[0005] Therefore, how to effectively reduce the computational complexity, storage space and power consumption of the model while ensuring its recognition accuracy has become an important direction in the research of intelligent inspection of power UAVs. SUMMARY

[0006] The purpose of the present application is to provide a defect recognition model lightweight method and system based on reinforcement learning to solve the problem of balancing the performance and lightweight requirements of deep learning models under the condition that the computing resources, storage space and battery capacity of edge devices are limited.

[0007] The present invention solves the above technical problems through the following technical solutions: a lightweight defect recognition model method based on reinforcement learning, comprising:

[0008] S1: Construct a sample data set according to the defect recognition model; wherein each sample in the sample data set includes a current state vector, a current action vector, a current instant reward, and a next state vector;

[0009] S2: Use the i-th sample to update the policy network and value network for the i-th time, and obtain the i-th updated policy network and value network;

[0010] S3: Prune the defect recognition model according to the updated policy network and the total target pruning ratio, and calculate the recognition accuracy of the pruned defect recognition model on the validation set;

[0011] S4: Based on the recognition accuracy and maximum number of iterations of the pruned defect recognition model on the validation set, determine whether the training of the policy network and the value network is completed. If so, output the trained policy network and proceed to S5; if not, proceed to S2 based on the i+1th sample;

[0012] S5: The defect recognition model is pruned according to the trained strategy network and the total target pruned ratio to obtain a lightweight defect recognition model.

[0013] Furthermore, in S1, the specific construction process of each sample in the sample data set includes:

[0014] Get random action vectors and total random cropping ratios;

[0015] The defect recognition model is clipped according to the random action vector and the total random clipping ratio to obtain a current state vector; wherein the current state vector is composed of the actual clipping ratios of each convolutional layer in the defect recognition model;

[0016] Using the unupdated policy network to perform feature extraction on the current state vector to obtain a current action vector;

[0017] Clipping the defect recognition model according to the current action vector and the total target clipping ratio to obtain a next state vector;

[0018] Calculate the recognition accuracy of the pruned defect recognition model on the validation set;

[0019] The current immediate reward is calculated based on the recognition accuracy of the pruned defect recognition model on the validation set and the current action vector.

[0020] Furthermore, the specific calculation formula for the current instant reward is:

[0021] ;

[0022] in, Indicates the current instant reward; represents the adjustment weight; represents the recognition accuracy of the pruned defect recognition model on the kth validation set; represents the j-th cropping ratio in the current action vector, that is, the cropping ratio of the j-th convolutional layer; M represents the number of validation sets; N represents the number of convolutional layers that can be cropped in the defect recognition model.

[0023] Furthermore, in S2, the policy network and the value network are updated for the i-th time using the i-th sample to obtain the i-th updated policy network and value network, which specifically includes:

[0024] Use the updated policy network after the i-1th time to extract features from the current state vector of the i-th sample and obtain the next action vector;

[0025] Use the updated value network after the i-1th time to extract features from the current state vector and current action vector of the i-th sample to obtain the current expected reward;

[0026] Use the value network after the i-1th update to extract the features of the next state vector and the next action vector of the i-th sample to obtain the next expected reward;

[0027] Calculate the loss function of the policy network after the i-1th update based on the current state vector, current action vector and current expected return of the i-th sample, and then perform the i-th update on the policy network after the i-1th update;

[0028] The loss function of the value network after the i-1th update is calculated based on the current expected return, current immediate reward and next expected return of the i-th sample, and then the value network is updated for the i-th time.

[0029] Furthermore, the calculation formula of the loss function of the policy network after the i-1th update is:

[0030] ;

[0031] in, Represents the loss function of the policy network after the i-1th update; represents the expected value calculation function; Represents a learnable parameter used to control the degree of policy exploration; Indicates the current state vector Output the current action vector probability; represents the current expected return;

[0032] ;

[0033] in, Represents the loss function of the value network after the i-1th update; Indicates the current instant reward; represents a learnable parameter that controls the impact of future rewards; Indicates the next expected return; represents the next state vector; represents the next action vector.

[0034] Furthermore, in S3, the defect recognition model is pruned according to the updated strategy network and the total target pruned ratio, specifically including:

[0035] Get random action vectors and total random cropping ratios;

[0036] Clipping the defect recognition model according to the random action vector and the total random clipping ratio to obtain a state vector;

[0037] Using the updated policy network for the i-th time to extract features from the state vector, obtaining an action vector;

[0038] The amount of allocated cropping parameters for each convolutional layer in the defect recognition model is calculated based on the action vector, the total target cropping ratio, and the total number of parameters of the defect recognition model. The specific calculation formula is:

[0039] , ;

[0040] ;

[0041] in, Indicates the total amount of trimmed parameters of the defect recognition model; Indicates the total target cropping ratio; Indicates the total number of parameters of the defect recognition model; represents the jth cropping ratio in the action vector, that is, the cropping ratio of the jth convolutional layer; N represents the number of convolutional layers that can be cropped in the defect recognition model; Represents the clipping weight coefficient of the j-th convolutional layer in the defect recognition model; Indicates the amount of allocated clipping parameters for the j-th convolutional layer in the defect recognition model;

[0042] For the jth convolutional layer, it is determined whether the total number of parameters of the convolutional layer is less than the allocated pruning parameter amount, and if so, a difference value between the allocated pruning parameter amount and the total number of parameters of the convolutional layer is calculated, the difference value is randomly allocated to other convolutional layers with a total number of parameters greater than the allocated pruning parameter amount, and the actual pruning parameter amount of each convolutional layer in the defect recognition model is obtained.

[0043] The defect recognition model is pruned according to the actual pruning parameter amount of each convolutional layer in the defect recognition model.

[0044] Further, the defect recognition model is pruned according to the actual pruning parameter amount of each convolutional layer in the defect recognition model, specifically including:

[0045] For each convolutional layer, the importance of each convolutional kernel is calculated according to the weight coefficient of each convolutional kernel in the convolutional layer.

[0046] According to the importance of each convolutional kernel, all convolutional kernels of the convolutional layer are sorted in descending order.

[0047] According to the convolutional kernel sorted in descending order and the actual pruning parameter amount of the convolutional layer, a candidate pruning set is constructed.

[0048] An optimal pruning convolutional kernel set is selected from the candidate pruning set using an optimization algorithm, and pruning is performed.

[0049] Further, the optimal pruning convolutional kernel set is selected from the candidate pruning set using a genetic algorithm, specifically including:

[0050] B1: A group of pruning convolutional kernels is randomly selected from the candidate pruning set according to the total target pruning ratio, the defect recognition model is pruned according to the group of pruning convolutional kernels, and the recognition accuracy of the pruned defect recognition model on the verification set is calculated, to obtain a candidate model and its recognition accuracy, and further obtain a candidate model set composed of each candidate model and its recognition accuracy.

[0051] B2: According to the recognition accuracy of each candidate model, a first candidate model and a second candidate model are selected from the candidate model set.

[0052] B3: Taking the recognition accuracy as the fitness function of the genetic algorithm, and taking the first candidate model and the second candidate model as the parents of the genetic algorithm, the pruning convolutional kernels corresponding to the first candidate model and the second candidate model are crossed and mutated to obtain a sub-candidate model and its recognition accuracy.

[0053] B4: It is determined whether the maximum number of iterations is reached, if not, the sub-candidate model and its recognition accuracy are added to the candidate model set, and B2 is entered; if yes, a group of convolutional kernels corresponding to the maximum recognition accuracy in the candidate model set is taken as the optimal pruning convolutional kernel set.

[0054] Based on the same concept, the present invention also provides a lightweight defect recognition model system based on reinforcement learning, comprising:

[0055] A construction unit, configured to construct a sample data set according to the defect recognition model; wherein each sample in the sample data set includes a current state vector, a current action vector, a current instant reward, and a next state vector;

[0056] An updating unit, configured to update the policy network and the value network for the i-th time using the i-th sample, and obtain the i-th updated policy network and value network;

[0057] The cropping and calculation unit is used to crop the defect recognition model according to the policy network after the i-th update and the total target cropping ratio, and calculate the recognition accuracy of the cropped defect recognition model on the validation set;

[0058] The judgment unit is used to determine whether the training of the policy network and the value network is completed based on the recognition accuracy and maximum number of iterations of the pruned defect recognition model on the validation set. If so, the trained policy network is output; if not, the update unit, the pruned and calculated unit, and the judgment unit are repeatedly called based on the i+1th sample;

[0059] The cropping unit is used to crop the defect recognition model according to the trained strategy network and the total target cropping ratio to obtain a lightweight defect recognition model.

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

[0061] The present invention outputs the cropping ratio of each convolutional layer of the defect recognition model through reinforcement learning, and crops the defect recognition model according to the cropping ratio of each convolutional layer and the total target cropping ratio. On the basis of ensuring the accuracy of defect recognition, the number of parameters of the defect recognition model is reduced, and the recognition performance and lightweight balance of the defect recognition model are achieved; the lightweight defect recognition model reduces the model's computational complexity, storage space and power consumption, so that the defect recognition model can be deployed on edge devices with limited resources (such as drones), thereby realizing real-time reasoning of defect recognition. BRIEF DESCRIPTION OF THE DRAWINGS

[0062] In order to more clearly illustrate the technical solution of the present invention, the following briefly introduces the drawings required for use in the description of the embodiments. Obviously, the drawings described below are only one embodiment of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying any creative work.

[0063] Figure 1 This is a flow chart of a lightweight defect recognition model method based on reinforcement learning in an embodiment of the present invention;

[0064] Figure 2 Schematic diagram of the process of constructing a candidate cropping set in an embodiment of the present invention;

[0065] Figure 3 is a schematic diagram of selecting the first candidate model and the second candidate model in an embodiment of the present invention;

[0066] Figure 4 Schematic diagram of crossover and mutation operations in an embodiment of the present invention. DETAILED DESCRIPTION

[0067] The following is a clear and complete description of the technical solutions of the present invention in conjunction with the accompanying drawings of the embodiments of the present invention. Obviously, the embodiments described are only part of the embodiments of the present invention, not all of the embodiments. All other embodiments obtained by ordinary technicians in this field based on the embodiments of the present invention without making any creative efforts are within the scope of protection of the present invention.

[0068] The following specific embodiments are used to describe the technical solution of the present application in detail. The following specific embodiments can be combined with each other, and the same or similar concepts or processes may not be described in detail in some embodiments.

[0069] Example 1

[0070] Figure 1 The flowchart of the lightweight defect recognition model based on reinforcement learning provided by the present invention is shown. Figure 1 As shown, the defect recognition model lightweight method includes the following steps:

[0071] S1: Construct a sample dataset based on the defect recognition model.

[0072] Taking power defect identification as an example, a deep neural network model is trained on a power inspection dataset to obtain a power defect identification model. The power defect identification model can be used to identify common defect types such as broken insulators, bird nests, and loose bolts. In this embodiment, YOLOv8 is used as the deep neural network model. The power inspection dataset contains power defect images. A training set and multiple validation sets can be constructed based on the power inspection dataset. The training set is used to train the deep neural network model, and the validation set is used to validate the trained deep neural network model. The number of input channels, number of output channels, convolution kernel size, and total number of parameters of the power defect identification model and its convolutional layers are obtained. At the same time, the recognition accuracy of the power defect identification model on each validation set is also obtained for subsequent calculation of instant rewards.

[0073] The lightweight process of achieving defect recognition models based on reinforcement learning essentially models the cropping strategy (or pruning strategy) as a sequential decision problem. Through interactive learning between reinforcement learning and the defect recognition model, the performance of the defect recognition model is maximized while reducing redundant convolution kernels. This method, based on the Soft Actor-Critic (SAC) algorithm in reinforcement learning, boasts excellent sample efficiency, strategy diversity, and training stability, making it particularly suitable for compression tasks such as cropping in continuous action spaces. Specifically, the compression task is first modeled as a Markov decision process, and its state vector, action vector, and reward function are defined. Each cropping operation is considered an interaction between reinforcement learning and the defect recognition model, with each state vector and corresponding action vector being a continuous vector.

[0074] In a specific embodiment of the present invention, the specific construction process of each sample in the sample data set includes:

[0075] S1.1: Get random action vectors and total random cropping ratios.

[0076] The random action vector consists of N random cropping ratios. Each random cropping ratio is randomly generated and corresponds to a convolutional layer of the defect recognition model. That is, the length of the random action vector N is equal to the number of convolutional layers in the defect recognition model. The total random cropping ratio is the total cropping ratio of the randomly generated defect recognition model.

[0077] S1.2: Crop the defect recognition model according to the random action vector and the total random cropping ratio in S1.1 to obtain the current state vector.

[0078] In a specific embodiment of the present invention, the defect recognition model is pruned according to the random action vector and the total random pruned ratio in S1.1, specifically including:

[0079] The amount of allocated cropping parameters for each convolutional layer in the defect recognition model is calculated based on the random action vector, the total random cropping ratio, and the total number of parameters of the defect recognition model. For specific calculations, see formulas (5) to (7);

[0080] For the jth convolutional layer, determine whether the total number of parameters of the convolutional layer is less than its allocated cropping parameter amount. If so, calculate the difference between the allocated cropping parameter amount of the convolutional layer and its total number of parameters, and randomly distribute the difference to other convolutional layers whose total number of parameters is greater than the allocated cropping parameter amount, and obtain the actual cropping parameter amount of each convolutional layer in the defect recognition model;

[0081] The defect recognition model is pruned according to the actual pruned parameter amount of each convolutional layer in the defect recognition model to obtain a current state vector; wherein the current state vector is composed of the actual pruned ratio of each convolutional layer in the defect recognition model (that is, determined by the actual pruned parameter amount).

[0082] S1.3: Use the unupdated policy network to extract features from the current state vector in S1.2 to obtain the current action vector.

[0083] The soft actor-critic algorithm includes a policy network and a value network. The policy network contains two fully connected hidden layers and an output terminal. The output terminal includes two parallel linear layers. , the policy network outputs the mean and logarithmic standard deviation respectively, and samples the corresponding actions from the standard normal distribution:

[0084] (1)

[0085] in, Represents the current action vector The j-th action in , that is, the cropping ratio of the j-th convolutional layer; Represents the current state vector The j-th state in , that is, the actual cropping ratio of the j-th convolutional layer; represents the mean; represents the standard deviation; represents the Hadamard product, that is, element-wise multiplication; is random noise sampled from the standard normal distribution N(0,1), which is used to introduce a certain degree of random exploration.

[0086] S1.4: Crop the defect recognition model according to the current action vector and the total target cropping ratio to obtain the next state vector.

[0087] In a specific embodiment of the present invention, the defect recognition model (i.e., the original model) is cropped according to the current action vector and the total target cropping ratio, specifically including:

[0088] Calculate the amount of allocated cropping parameters for each convolutional layer in the defect recognition model based on the current action vector, the total target cropping ratio, and the total number of parameters of the defect recognition model;

[0089] For the jth convolutional layer, determine whether the total number of parameters of the convolutional layer is less than its allocated cropping parameter amount. If so, calculate the difference between the allocated cropping parameter amount of the convolutional layer and its total number of parameters, and randomly distribute the difference to other convolutional layers whose total number of parameters is greater than the allocated cropping parameter amount, and obtain the actual cropping parameter amount of each convolutional layer in the defect recognition model;

[0090] The defect recognition model is trimmed according to the actual trimming parameters of each convolutional layer in the defect recognition model to obtain the next state vector .

[0091] S1.5: Calculate the recognition accuracy of the pruned defect recognition model on the validation set.

[0092] The images in the validation set are input into the cropped defect recognition model to obtain the defect recognition results; the recognition accuracy is calculated based on the defect recognition results and the actual defect type, thereby obtaining the recognition accuracy of the cropped defect recognition model on each validation set.

[0093] S1.6: Calculate the current instant reward based on the recognition accuracy of the pruned defect recognition model on the validation set and the current action vector. The specific calculation formula is:

[0094] (2)

[0095] in, Indicates the current instant reward; Represents the adjustment weight, which is used to adjust the ratio of the precision part to the clipping part in the immediate reward function. In this embodiment, it is set to 0.8, indicating a bias towards clipping precision; Represents the recognition accuracy of the defect recognition model on the kth validation set; represents the j-th cropping ratio in the current action vector, that is, the cropping ratio of the j-th convolutional layer; M represents the number of validation sets; N represents the number of convolutional layers that can be cropped in the defect recognition model.

[0096] During each reinforcement learning session, the policy network generates an action vector consisting of cropping ratios, which is applied to the original defect recognition model to generate a cropped defect recognition model. The cropped defect recognition model is then evaluated for accuracy, with the resulting recognition accuracy index used as a reward. Since a convolutional layer with more convolution kernels has a greater cropping potential, in order to maximize the overall cropping effect of the model and avoid the performance waste caused by uniform cropping, the present invention introduces the reciprocal sum of the cropping ratios as an auxiliary term in the instantaneous reward calculation. The auxiliary term contributes more to convolutional layers with a larger number of convolution kernels, and therefore can effectively guide the policy network to perform greater cropping in convolutional layers with a larger number of convolution kernels.

[0097] Therefore, each sample in the sample dataset includes the current state vector , current action vector 、Current instant rewards and the next state vector .

[0098] S2: Use the i-th sample to update the policy network and value network for the i-th time, and obtain the i-th updated policy network and value network.

[0099] In a specific embodiment of the present invention, the policy network and the value network are updated for the i-th time using the i-th sample to obtain the i-th updated policy network and value network, which specifically includes:

[0100] S2.1: Use the updated policy network after the i-1th update to extract features from the current state vector of the i-th sample and obtain the next action vector.

[0101] For the current state vector of the i-th sample of input , the policy network outputs the mean and logarithmic standard deviation respectively, and samples the corresponding action from the standard normal distribution, as shown in formula (1), and then obtains the next action vector .

[0102] S2.2: Use the updated value network to extract features from the current state vector and current action vector of the i-th sample to obtain the current expected reward.

[0103] The value network adopts a dual critic network structure. Each critic network takes the concatenation of the current state vector and the current action vector as input and predicts the expected return of the current action vector under the current state vector, that is, the current expected return .

[0104] S2.3: Use the updated value network after the i-1th time to extract the next state vector and next action vector of the i-th sample to obtain the next expected reward .

[0105] S2.4: Calculate the loss function of the policy network after the i-1th update based on the current state vector, current action vector, and current expected reward of the i-th sample, and then perform the i-th update on the policy network after the i-1th update.

[0106] The goal of the policy network is to learn a probability distribution , in the state vector Choose the action that will bring the greatest long-term benefits , long-term benefits are measured by the Q value (i.e., expected return) output by the value network. Therefore, the policy network tends to choose actions with larger Q values. In addition, the policy network also introduces policy entropy to measure the uncertainty or diversity of the policy, encouraging the exploration of the policy. Therefore, the loss function of the policy network is:

[0107] (3)

[0108] in, Represents the loss function of the policy network after the i-1th update; represents the expected value calculation function; Represents a learnable parameter used to control the degree of policy exploration; Indicates the current state vector Output the current action vector The probability of This is the policy entropy term; Represents the current expected return, which is approximated to its true value through the value network.

[0109] S2.5: Calculate the loss function of the value network after the i-1th update based on the current expected return, current immediate reward, and next expected return of the i-th sample, and then update the value network for the i-th time.

[0110] In this embodiment, the loss function of the value network is:

[0111] (4)

[0112] in, Represents the loss function of the value network after the i-1th update; Indicates the current instant reward; represents a learnable parameter that controls the impact of future rewards, The value range is 0~1; Indicates the next expected return; represents the next state vector; represents the next action vector. Represents the actual reward and the estimated future reward, reflecting the balance between current and future gains, and is used to guide the approximation of the correct action value.

[0113] S3: Prune the defect recognition model according to the updated policy network and the total target pruning ratio, and calculate the recognition accuracy of the pruned defect recognition model on the validation set.

[0114] In a specific embodiment of the present invention, the defect recognition model is pruned according to the strategy network after the i-th update and the total target pruned ratio, specifically including:

[0115] S3.1: Get random action vectors and total random cropping ratios.

[0116] S3.2: Prune the defect recognition model according to the random action vector and the total random pruning ratio in S3.1 to obtain the state vector. For the specific pruning process, see S1.2.

[0117] S3.3: Use the updated policy network to extract features from the state vector obtained in S3.2 to obtain the action vector.

[0118] S3.4: Calculate the amount of allocated cropping parameters for each convolutional layer in the defect recognition model based on the action vector, the total target cropping ratio, and the total number of parameters of the defect recognition model. The specific calculation formula is:

[0119] (5)

[0120] (6)

[0121] (7)

[0122] in, Indicates the total amount of trimmed parameters of the defect recognition model; Indicates the total target cropping ratio; Indicates the total number of parameters of the defect recognition model; represents the jth cropping ratio in the action vector, that is, the cropping ratio of the jth convolutional layer; N represents the number of convolutional layers that can be cropped in the defect recognition model; Represents the clipping weight coefficient of the j-th convolutional layer in the defect recognition model; Represents the amount of allocated pruning parameters for the j-th convolutional layer in the defect recognition model.

[0123] S3.5: For the j-th convolutional layer, determine whether the total number of parameters of the convolutional layer is less than the number of parameters allocated for clipping;

[0124] If not, the actual amount of cropping parameters of the convolutional layer is equal to its allocated amount of cropping parameters;

[0125] If so, the difference between the allocated cropping parameter amount of the convolutional layer and its total parameter amount is calculated, and the difference is randomly distributed to other convolutional layers whose total parameter amount is greater than the allocated cropping parameter amount, so as to obtain the actual cropping parameter amount of each convolutional layer in the defect recognition model.

[0126] S3.5 ensures that the actual pruned parameters of each convolutional layer do not exceed its total parameters.

[0127] S3.6: Prune the defect recognition model based on the actual number of pruned parameters of each convolutional layer in the defect recognition model.

[0128] In a specific embodiment of the present invention, the defect recognition model is trimmed according to the actual trimming parameter amount of each convolutional layer in the defect recognition model, specifically including:

[0129] S3.61: For each convolutional layer, the importance of each convolution kernel is calculated based on the weight coefficient of each convolution kernel in the convolutional layer. The specific calculation formula is:

[0130] (8)

[0131] in, Represents the weight tensor of the convolution kernel in the convolution layer; Represents the square of the L2 norm of the weight tensor of the convolution kernel; Represents the convolution weight of the kth height and lth width on the i-th output channel and the j-th input channel, which is the weight parameter of the defect recognition model; Indicates the number of output channels of the convolutional layer; Indicates the number of input channels of the convolutional layer; Indicates the height of the convolution kernel; Indicates the width of the convolution kernel.

[0132] S3.62: Sort all convolution kernels of this convolution layer in descending order according to the importance of each convolution kernel.

[0133] S3.63: Construct a candidate cropping set based on the descending sorted convolution kernels and the actual cropping parameters of the convolution layer.

[0134] After descending sorting, the convolution kernels with the lowest ranking are added to the candidate cropping set based on the total target cropping ratio. To ensure that the candidate cropping set is not entirely determined by the importance index of the convolution kernel and to increase the diversity of the candidate cropping set, the present invention first adds the convolution kernel with the lowest total target cropping ratio to the candidate cropping set, and then adds some convolution kernels before this to the candidate cropping set to expand the candidate cropping set.

[0135] For example, if the total target cropping ratio is 40%, and the first ranking is 100% and the last ranking is 0, the convolution kernels ranked from 40% to 0 are first added to the candidate cropping set, and then the convolution kernels ranked from 60% to 40% are added to the candidate cropping set. Figure 2 The figure shows the process of constructing the candidate crop set. The top-ranked convolution kernels are directly retained, the last-ranked convolution kernels are directly added to the candidate crop set, and the middle part is used to expand the candidate crop set.

[0136] S3.64: Use the optimization algorithm to select the optimal cropping convolution kernel set from the candidate cropping set and perform cropping.

[0137] In a specific embodiment of the present invention, a genetic algorithm is used to select the optimal cropping convolution kernel set from the candidate cropping set, specifically including:

[0138] B1: randomly select a set of pruning convolution kernels from the candidate pruning set according to the total target pruning ratio, prune the defect recognition model according to the set of pruning convolution kernels, calculate the recognition accuracy of the pruned defect recognition model on the verification set, obtain the candidate model and the recognition accuracy, and further obtain a candidate model set composed of each candidate model and the recognition accuracy, as shown in Figure 3 The pruning convolution kernel refers to the convolution kernel to be pruned.

[0139] B2: select a first candidate model and a second candidate model from the candidate model set according to the recognition accuracy of each candidate model.

[0140] In this embodiment, a selection probability distribution is constructed according to the recognition accuracy of each candidate model in the candidate model set, and the first candidate model and the second candidate model are selected based on the selection probability distribution, that is, the higher the recognition accuracy of the candidate model, the greater the probability of being selected. Wherein, the recognition accuracy of each candidate model is normalized into a selection probability, and the specific formula is:

[0141] (9)

[0142] Wherein, represents the selection probability of the kth candidate model; represents the recognition accuracy of the kth candidate model; represents the number of candidate models in the candidate model set.

[0143] B3: taking the recognition accuracy as the fitness function of the genetic algorithm, and taking the first candidate model and the second candidate model as the parents of the genetic algorithm, the pruning convolution kernels corresponding to the first candidate model and the second candidate model are crossed and mutated to obtain a sub candidate model and its recognition accuracy, as shown in Figure 4 .

[0144] Crossing: sequentially traversing each convolution layer of the first candidate model and the second candidate model, randomly exchanging the pruning convolution kernel corresponding to the convolution layer of the first candidate model and the pruning convolution kernel corresponding to the convolution layer of the second candidate model to generate a new pruning convolution kernel corresponding to the convolution layer, pruning the defect recognition model according to the new pruning convolution kernel corresponding to all convolution layers and calculating the recognition accuracy of the pruned defect recognition model on the verification set, obtaining an initial sub candidate model and its recognition accuracy;

[0145] Mutation: Traverse each convolutional layer of the first candidate model and the second candidate model in turn, select several cropped convolution kernels from the cropped convolution kernels corresponding to the convolutional layer of the first candidate model and the cropped convolution kernels corresponding to the convolutional layer of the second candidate model, respectively, to replace the cropped convolution kernels of the corresponding convolutional layer of the initial sub-candidate model, crop the defect recognition model according to the cropped convolution kernels obtained after replacement corresponding to all convolution layers, and calculate the recognition accuracy of the cropped defect recognition model on the validation set to obtain the sub-candidate model and its recognition accuracy.

[0146] B4: Determine whether the maximum number of iterations has been reached. If not, add the sub-candidate model and its recognition accuracy to the candidate model set and proceed to B2. If so, use the set of convolution kernels corresponding to the maximum recognition accuracy in the candidate model set as the optimal cropped convolution kernel set.

[0147] The present invention introduces a genetic algorithm to search for the optimal set of cropping convolution kernels (i.e., the optimal cropping position). By increasing the diversity of selection through crossover and mutation operations, it helps to escape the local optimum and expand the search space. By evaluating the fitness of new individuals through recognition accuracy feedback, the present invention realizes the fusion of global search and local fine-tuning of the cropping position, which not only ensures the coverage of a wide solution space, but also enables the fine-tuning of the optimal solution.

[0148] S4: Based on the recognition accuracy and maximum number of iterations of the pruned defect recognition model on the validation set, determine whether the training of the policy network and the value network is completed. If so, output the trained policy network and proceed to S5; if not, proceed to S2 based on the i+1th sample.

[0149] If the maximum number of iterations has not been reached and the recognition accuracy of the pruned defect recognition model on the validation set reaches the accuracy threshold, the training of the policy network and the value network is considered complete, and a trained policy network is obtained. If the maximum number of iterations has been reached and the recognition accuracy of the pruned defect recognition model on the validation set has not reached the accuracy threshold, the training of the policy network and the value network is considered complete, and a trained policy network is obtained. In other words, training is considered complete as long as either the recognition accuracy of the pruned defect recognition model on the validation set reaches the accuracy threshold or the maximum number of iterations has been reached.

[0150] S5: The defect recognition model is pruned according to the trained strategy network and the total target pruned ratio to obtain a lightweight defect recognition model.

[0151] S5 is the application deployment of the trained policy network. In a specific embodiment of the present invention, the defect recognition model is tailored according to the trained policy network and the total target tailoring ratio, specifically including:

[0152] S5.1: Get random action vectors and total random cropping ratios.

[0153] S5.2: According to the random action vector in S5.1 and the total random pruning ratio, the defect recognition model is pruned to obtain a state vector.

[0154] S5.3: The state vector obtained in S5.2 is subjected to feature extraction by using the trained strategy network to obtain an action vector.

[0155] S5.4: According to the action vector, the total target pruning ratio and the total parameter amount of the defect recognition model, the allocated pruning parameter amount of each convolutional layer in the defect recognition model is calculated.

[0156] S5.5: For the jth convolutional layer, it is judged whether the total parameter amount of the convolutional layer is less than the allocated pruning parameter amount thereof;

[0157] If not, the actual pruning parameter amount of the convolutional layer is equal to the allocated pruning parameter amount thereof;

[0158] If yes, the difference between the allocated pruning parameter amount and the total parameter amount of the convolutional layer is calculated, and the difference is randomly allocated to other convolutional layers with a total parameter amount greater than the allocated pruning parameter amount, to obtain the actual pruning parameter amount of each convolutional layer in the defect recognition model.

[0159] S5.6: According to the actual pruning parameter amount of each convolutional layer in the defect recognition model, the defect recognition model is pruned, i.e. a lightweight defect recognition model is obtained.

[0160] The present application constructs a pruning intelligent agent (realized by a strategy network-value network) capable of continuously adjusting the pruning ratio of the defect recognition model according to the performance of the pruned defect recognition model, realizes differential pruning of different convolutional layers of the defect recognition model, and continuously explores the position of the pruned convolutional kernel through the importance index of the convolutional kernel and the optimization algorithm, while maintaining the original performance of the defect recognition model to the greatest extent.

[0161] Specifically, the pruning intelligent agent generates a pruning action, i.e. a pruning ratio, by using a strategy network, simulates the behavior of the model after pruning operation through an environment, and simultaneously feeds back the recognition accuracy of the pruned defect recognition model on the validation set as the main reward signal of the immediate reward, and combines the action vector to guide the pruning intelligent agent to prune more in the convolutional layer with a large number of convolutional kernels. This immediate reward mechanism enables the pruning intelligent agent to automatically learn how to maintain the recognition accuracy to the greatest extent while pruning the structure of the defect recognition model, thereby generating a reasonable structure pruning strategy.

[0162] The training process of the strategy network-value network is the interaction process of the pruning intelligent agent and the environment, and a single sample (a defect recognition model) ) is also obtained by cropping the interaction between the agent and the environment, storing multiple samples as experience in the Replay Buffer, and randomly sampling a batch of samples from the Replay Buffer during training to update the policy network and value network to obtain a trained policy network.

[0163] Example 2

[0164] The embodiment of the present invention also provides a lightweight defect recognition model system based on reinforcement learning, which includes a construction unit, an update unit, a clipping and calculation unit, a judgment unit, and a clipping unit.

[0165] A construction unit, configured to construct a sample data set according to the defect recognition model; wherein each sample in the sample data set includes a current state vector, a current action vector, a current instant reward, and a next state vector;

[0166] An updating unit, configured to update the policy network and the value network for the i-th time using the i-th sample, and obtain the i-th updated policy network and value network;

[0167] The cropping and calculation unit is used to crop the defect recognition model according to the policy network after the i-th update and the total target cropping ratio, and calculate the recognition accuracy of the cropped defect recognition model on the validation set;

[0168] The judgment unit is used to determine whether the training of the policy network and the value network is completed based on the recognition accuracy and maximum number of iterations of the pruned defect recognition model on the validation set. If so, the trained policy network is output; if not, the update unit, the pruned and calculated unit, and the judgment unit are repeatedly called based on the i+1th sample;

[0169] The cropping unit is used to crop the defect recognition model according to the trained strategy network and the total target cropping ratio to obtain a lightweight defect recognition model.

[0170] In some specific embodiments of the present invention, the defect recognition model lightweight system may be combined with the features of the defect recognition model lightweight method in the first embodiment of the present invention, and vice versa.

[0171] The above disclosure is only a specific embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any technician familiar with this technical field can easily think of changes or modifications within the technical scope disclosed in the present invention, and they should all be covered by the scope of protection of the present invention.

Claims

1. A lightweight defect recognition model method based on reinforcement learning, characterized in that: The lightweight method includes: S1: Construct a sample data set according to the defect recognition model; wherein each sample in the sample data set includes a current state vector, a current action vector, a current instant reward, and a next state vector; S2: Use the i-th sample to update the policy network and value network for the i-th time, and obtain the i-th updated policy network and value network; S3: Prune the defect recognition model according to the updated policy network and the total target pruning ratio, and calculate the recognition accuracy of the pruned defect recognition model on the validation set; S4: Based on the recognition accuracy and maximum number of iterations of the pruned defect recognition model on the validation set, determine whether the training of the policy network and the value network is completed. If so, output the trained policy network and proceed to S5; if not, proceed to S2 based on the i+1th sample; S5: The defect recognition model is pruned according to the trained strategy network and the total target pruned ratio to obtain a lightweight defect recognition model. In S3, the defect recognition model is pruned according to the updated policy network and the total target pruned ratio, specifically including: Get random action vectors and total random cropping ratios; Clipping the defect recognition model according to the random action vector and the total random clipping ratio to obtain a state vector; Using the updated policy network for the i-th time to extract features from the state vector, obtaining an action vector; The amount of allocated cropping parameters for each convolutional layer in the defect recognition model is calculated based on the action vector, the total target cropping ratio, and the total number of parameters of the defect recognition model. The specific calculation formula is: ， ； ； in, Indicates the total amount of trimmed parameters of the defect recognition model; Indicates the total target cropping ratio; Represents the total number of parameters of the defect recognition model; represents the jth cropping ratio in the action vector, that is, the cropping ratio of the jth convolutional layer; N represents the number of convolutional layers that can be cropped in the defect recognition model; Represents the clipping weight coefficient of the j-th convolutional layer in the defect recognition model; Indicates the amount of allocated clipping parameters for the jth convolutional layer in the defect recognition model; For the jth convolutional layer, determine whether the total number of parameters of the convolutional layer is less than its allocated cropping parameter amount. If so, calculate the difference between the allocated cropping parameter amount of the convolutional layer and its total number of parameters, and randomly distribute the difference to other convolutional layers whose total number of parameters is greater than the allocated cropping parameter amount, and obtain the actual cropping parameter amount of each convolutional layer in the defect recognition model; The defect recognition model is pruned according to the actual pruned parameter amount of each convolutional layer in the defect recognition model.

2. The lightweight defect recognition model method based on reinforcement learning according to claim 1 is characterized in that: In S1, the specific construction process of each sample in the sample data set includes: Get random action vectors and total random cropping ratios; The defect recognition model is clipped according to the random action vector and the total random clipping ratio to obtain a current state vector; wherein the current state vector is composed of the actual clipping ratios of each convolutional layer in the defect recognition model; Using the unupdated policy network to perform feature extraction on the current state vector to obtain a current action vector; Clipping the defect recognition model according to the current action vector and the total target clipping ratio to obtain a next state vector; Calculate the recognition accuracy of the pruned defect recognition model on the validation set; The current immediate reward is calculated based on the recognition accuracy of the pruned defect recognition model on the validation set and the current action vector.

3. The lightweight defect recognition model method based on reinforcement learning according to claim 2 is characterized in that: The specific calculation formula for the current instant reward is: ； in, Indicates the current instant reward; represents the adjustment weight; represents the recognition accuracy of the pruned defect recognition model on the kth validation set; represents the j-th cropping ratio in the current action vector, that is, the cropping ratio of the j-th convolutional layer; M represents the number of validation sets; N represents the number of convolutional layers that can be cropped in the defect recognition model.

4. The lightweight defect recognition model method based on reinforcement learning according to claim 1 is characterized in that: In S2, the policy network and the value network are updated for the i-th time using the i-th sample to obtain the i-th updated policy network and value network, which specifically includes: Use the updated policy network after the i-1th time to extract features from the current state vector of the i-th sample and obtain the next action vector; Use the updated value network (i-1 times) to extract features from the current state vector and current action vector of the i-th sample to obtain the current expected reward. Use the value network after the i-1th update to extract the next state vector and next action vector of the i-th sample to obtain the next expected reward; Calculate the loss function of the policy network after the i-1th update based on the current state vector, current action vector and current expected return of the i-th sample, and then perform the i-th update on the policy network after the i-1th update; The loss function of the value network after the i-1th update is calculated based on the current expected return, current immediate reward and next expected return of the i-th sample, and then the value network is updated for the i-th time.

5. The lightweight defect recognition model method based on reinforcement learning according to claim 4 is characterized in that: The calculation formula of the loss function of the policy network after the i-1th update is: ； in, Represents the loss function of the policy network after the i-1th update; represents the expected value calculation function; Represents a learnable parameter used to control the degree of policy exploration; Indicates the current state vector Output the current action vector probability; represents the current expected return; ； in, Represents the loss function of the value network after the i-1th update; Indicates the current instant reward; represents a learnable parameter that controls the impact of future rewards; Indicates the next expected return; represents the next state vector; represents the next action vector.

6. The lightweight defect recognition model method based on reinforcement learning according to claim 1 is characterized in that: The defect recognition model is pruned based on the actual pruned parameters of each convolutional layer in the defect recognition model, specifically including: For each convolution layer, the importance of each convolution kernel is calculated according to the weight coefficient of each convolution kernel in the convolution layer; According to the importance of each convolution kernel, all convolution kernels of the convolution layer are sorted in descending order; Constructing a candidate cropping set according to the descending sorted convolution kernels and the actual cropping parameters of the convolution layer; An optimization algorithm is used to select the optimal cropping convolution kernel set from the candidate cropping set and perform cropping.

7. The lightweight defect recognition model method based on reinforcement learning according to claim 6 is characterized in that: Selecting the optimal cropping convolution kernel set from the candidate cropping set using a genetic algorithm, specifically including: B1: Randomly select a set of cropping convolution kernels from the candidate cropping set according to the total target cropping ratio, crop the defect recognition model according to the set of cropping convolution kernels, and calculate the recognition accuracy of the cropped defect recognition model on the validation set to obtain the candidate model and its recognition accuracy, and then obtain a candidate model set consisting of the candidate models and their recognition accuracy; B2: Selecting a first candidate model and a second candidate model from the candidate model set according to the recognition accuracy of each candidate model; B3: Using recognition accuracy as the fitness function of the genetic algorithm, and the first and second candidate models as the parents of the genetic algorithm, crossover and mutation are performed on the cropped convolution kernels corresponding to the first and second candidate models to obtain the child candidate models and their recognition accuracy; B4: Determine whether the maximum number of iterations has been reached. If not, add the sub-candidate model and its recognition accuracy to the candidate model set and proceed to B2. If so, use the set of convolution kernels corresponding to the maximum recognition accuracy in the candidate model set as the optimal cropped convolution kernel set.

8. A lightweight defect recognition model system based on reinforcement learning, characterized in that: The lightweight system includes: A construction unit, configured to construct a sample data set according to the defect recognition model; wherein each sample in the sample data set includes a current state vector, a current action vector, a current instant reward, and a next state vector; An updating unit, configured to update the policy network and the value network for the i-th time using the i-th sample, and obtain the i-th updated policy network and value network; The cropping and calculation unit is used to crop the defect recognition model according to the policy network after the i-th update and the total target cropping ratio, and calculate the recognition accuracy of the cropped defect recognition model on the validation set; The judgment unit is used to determine whether the training of the policy network and the value network is completed based on the recognition accuracy and maximum number of iterations of the pruned defect recognition model on the validation set. If so, the trained policy network is output; if not, the update unit, the pruned and calculated unit, and the judgment unit are repeatedly called based on the i+1th sample; The cropping unit is used to crop the defect recognition model according to the trained strategy network and the total target cropping ratio to obtain a lightweight defect recognition model; The clipping and calculation unit is used to clip the defect recognition model according to the strategy network after the i-th update and the total target clipping ratio, specifically including: Get random action vectors and total random cropping ratios; Clipping the defect recognition model according to the random action vector and the total random clipping ratio to obtain a state vector; Using the updated policy network for the i-th time to extract features from the state vector, obtaining an action vector; The amount of allocated cropping parameters for each convolutional layer in the defect recognition model is calculated based on the action vector, the total target cropping ratio, and the total number of parameters of the defect recognition model. The specific calculation formula is: ， ； ； in, Indicates the total amount of trimmed parameters of the defect recognition model; Indicates the total target cropping ratio; Indicates the total number of parameters of the defect recognition model; represents the jth cropping ratio in the action vector, that is, the cropping ratio of the jth convolutional layer; N represents the number of convolutional layers that can be cropped in the defect recognition model; Represents the clipping weight coefficient of the j-th convolutional layer in the defect recognition model; Indicates the amount of allocated clipping parameters for the jth convolutional layer in the defect recognition model; For the jth convolutional layer, determine whether the total number of parameters of the convolutional layer is less than its allocated cropping parameter amount. If so, calculate the difference between the allocated cropping parameter amount of the convolutional layer and its total number of parameters, and randomly distribute the difference to other convolutional layers whose total number of parameters is greater than the allocated cropping parameter amount, and obtain the actual cropping parameter amount of each convolutional layer in the defect recognition model; The defect recognition model is pruned according to the actual pruned parameter amount of each convolutional layer in the defect recognition model.

Images