CN116306813A - Method based on YOLOX light weight and network optimization - Google Patents

Abstract

The invention discloses a method for lightweight and network optimization based on YOLOX, which comprises the following steps of; s1: in the target detection task, preparing a data set required during training; s2: training an original YOLOX neural network model on the data set, and recording and evaluating performance indexes of the model; s3: performing pruning operation on the original YOLOX neural network model to generate a pruned improved YOLOX neural network model; s4: training on the dataset to generate a pruned improved YOLOX network model; s5: performing pruning operations on the improved YOLOX network model; s6: verifying and analyzing the improved YOLOX network after improved pruning, and if the requirements on performance can be met, detecting and analyzing the target; if the performance requirements cannot be met, the improvement model is adjusted until the performance requirements are met. The invention has higher detection precision and speed in target detection, is easier to deploy and integrate in actual application scenes, and also enables the reasoning process of the model to be more efficient and stable.

Description

A method based on YOLOX lightweight and network optimization

technical field

The invention belongs to the technical field of object detection in images, and in particular relates to a YOLOX-based lightweight and network optimization method.

Background technique

The object detection problem is to determine the location of objects in a given image and the class to which each object belongs (i.e., object localization and object classification). Nowadays, target detection technology has been widely used in agriculture, medical treatment, and automated production. Target detection mainly adopts the method of deep learning, and the target detection methods based on deep learning are mainly divided into two categories. One is a two-stage target detection algorithm based on candidate regions, including R-CNN, SPP-Net, Fast-RCNN, Faster-RCNN, etc.

Since these two-stage algorithms need to generate a large number of candidate regions, the target detection takes a relatively long time. It needs to first generate a large number of candidate regions, and then perform target classification and positioning on these candidate regions. This process consumes a lot of computing resources and time, especially in high-resolution images and complex scenes, which is relatively slow. Compared with the two-stage algorithm, the single-stage algorithm only needs to densely sample the image, and then directly perform target detection on the sampling points, so it does not need to generate a large number of candidate regions, the detection speed is faster, and it is suitable for real-time target detection. Therefore, the single-stage algorithm is more popular in practical applications.

As an excellent single-target detection algorithm, YOLOX can return the probability and position coordinates of object categories, and the algorithm is fast. However, when directly using YOLOX to detect objects, there will inevitably be some problems. The YOLOX model is large and the reasoning speed is slow, so it needs to be lightweight. While the YOLOX network is lightweight, it is necessary to maintain the original accuracy without loss. Therefore, it is necessary to maintain its detection accuracy through network optimization after lightweighting. .

Contents of the invention

In order to overcome the deficiencies of the above-mentioned prior art, the object of the present invention is to provide a method based on YOLOX lightweight and network optimization, which solves the limitations of model deployment in target detection and the problem of decreased accuracy in model performance evaluation after lightweight It has high detection accuracy and speed in target detection, is easier to deploy and integrate in actual application scenarios, and also makes the model's reasoning process more efficient and stable.

In order to achieve the above object, the technical scheme adopted in the present invention is:

A method based on YOLOX lightweight and network optimization, comprising the following steps;

S1: In the target detection task, prepare the data set required for training; the data set selects the original image in different scenes and different lighting conditions, and preprocesses the original image to ensure that the model can be used in various environments achieve accurate target detection;

S2: Train the original YOLOX neural network model on the data set, record and evaluate the performance indicators of the model;

S3: Perform a pruning operation on the original YOLOX neural network model to generate an improved YOLOX network model after pruning;

S4: Training and generating the improved YOLOX network model after pruning on the dataset;

S5: Perform a pruning operation on the improved YOLOX network model, and make further adjustments to achieve higher detection accuracy and speed;

S6: Verify and analyze the improved YOLOX network after pruning. If the performance requirements can be met, the target will be detected and analyzed; if the performance requirements cannot be met, the improved model will be adjusted until the performance requirements are met.

In the step S1, preparing the required data set for training includes the following steps:

(1) Collect data: collect PASCAL VOC datasets from the official website of open datasets, including JPEGImages, ImageSets and Annotations; where JPEGImages contains training datasets, and ImageSets contains each type of train.txt, trainval.txt and val. txt file, Annotations contains each type of xml file;

(2) Data preprocessing: Data preprocessing is to preprocess the collected original images to make them suitable for model training; first, the original images are adjusted to the specified size for subsequent processing, and the selected target detection original images The size is 416x416 pixels; next, the color image is converted to a grayscale image, which reduces the complexity of data storage and processing, and can reduce the time and calculation of model training. Finally, the pixel value in the image is scaled to 0 to 1 Between, making the data more stable in the process of processing, and can reduce the problem of gradient explosion and gradient disappearance in the training process.

Said step S2 trains the original YOLOX neural network model on the data set, comprising the following steps:

(1) adopt the data set after step S1 preprocessing;

(2) Divide the data set into training set and verification set according to the ratio of 8:2;

(3) Input the preprocessed data set into the original YOLOX neural network model, input the predicted value pred and the real value gt of the network into the loss function L, and calculate the loss value by the following formula

Loss=L(pred,gt)

Among them, L represents the loss function, pred represents the predicted value of the network output, gt represents the real value, optimize the network parameters according to the loss function L, use the gradient descent method to update the neural network parameters, set the current neural network parameters to θ, then update the formula for:

表示损失函数L对参数θ的梯度，θ表示第t个时间步的参数值，θ_t+1表示第t+1个时间步的参数值，通过多次迭代更新神经网络参数，优化网络性能，提高目标检测的准确率和速度；Among them, η represents the learning rate,

Represents the gradient of the loss function L to the parameter θ, θ represents the parameter value of the t-th time step, θ _t+1 represents the parameter value of the t+1-th time step, and updates the neural network parameters through multiple iterations to optimize network performance. Improve the accuracy and speed of target detection;

(4) After a round of parameter updating, the model needs to be tested using the verification set to verify the generalization ability of the model. Specifically, the verification set is input into the YOLOX network to calculate the loss between the predicted result and the real result The measurement, that is, the loss degree of the verification set, assuming that the size of the verification set is N, the prediction frame of the i-th sample is p _i , and the real frame is t _i , then the loss degree L of the verification set can be calculated as follows:

和/>

分别表示第i个样本的第j个格子中第c个类别的预测值和真实值，/>

和/>

分别表示第i个样本的第j个格子是否存在目标的预测值和真实值，/>

和/>

分别表示第i个样本的第j个格子中的置信度预测值和真实值，pos_ij表示第i个样本的第j个格子中与真实框有最大交并比的预测框的索引集合/>

和/>

是两个权重系数，用于平衡存在目标的格子和不存在目标的格子的权重；Among them, S is the number of prediction boxes for each grid, C is the number of target categories,

and />

respectively represent the predicted value and actual value of the c-th category in the j-th grid of the i-th sample, />

and />

Respectively indicate whether there is a predicted value and a real value of the target in the j-th grid of the i-th sample, />

and />

respectively represent the confidence prediction value and the real value in the jth grid of the i-th sample, pos _ij represents the index set of the prediction frame with the largest intersection and union ratio with the real frame in the j-th grid of the i-th sample />

and />

are two weight coefficients, which are used to balance the weights of grids with targets and grids without targets;

By calculating the loss degree of the verification set, the performance of the current model can be evaluated. If the loss degree is high, the training needs to be continued until the predetermined stop condition is reached;

(5) Every iteration twice, input the pictures in the data set into the optimized YOLOX target detection network for training to obtain the accuracy of the model. The optimized YOLOX target detection network has higher detection accuracy and faster detection speed;

(6) Repeat the above steps until the training ends.

In the step S3, the pruning operation is performed on the network trained in the step S2, including the following steps:

(1) According to the importance index of the network layer, determine the layer to be pruned; for each network layer, calculate the sensitivity of a network layer on the forward propagation of the model, and calculate the output change under the given input Quantify the partial derivative of the weight of the layer to obtain the sensitivity of the layer to the output. The greater the sensitivity, the greater the impact of the layer on the output, which should be given priority when pruning;

(2) Sorting the weights of the network layers to be pruned;

(3) Determine the threshold according to the weight sorting result and the pruning rate in each pruning layer;

(4) Eliminate weights below the threshold in the network and retain weights above the threshold;

(5) Save the new model parameters and weights, and generate the improved YOLOX network model after pruning.

Described step S4 improves YOLOX network model, comprises the following steps:

(1) Insert the channel-spatial attention mechanism CBAM module between the yolox backbone layer and the data enhancement layer channel;

The CBAM module is an implementation of the channel-spatial attention mechanism, which can effectively improve the accuracy of the model; it mainly includes two parts: channel attention and spatial attention;

First, through the channel attention mechanism, the input feature map is subjected to average pooling and maximum pooling operations respectively to realize the spatial information of the aggregated feature map, and the generated average pooling feature F _avg and maximum pooling feature F _max pass through the shared network layer , after applying the shared network to each feature, the average pooled features and the maximum pooled features are element-wise summed, and the combined features are output through the Sigmiod activation function. The channel attention map Mc; the spatial attention is aligned with the feature along the channel axis The average pooling and maximum pooling operations are performed on the graph, so that the feature map is compressed in the channel dimension, and the two feature maps are spliced in the channel dimension to generate an effective feature map, and then go through a 7X7 convolutional layer; finally, Obtain the final channel attention map Ms through the Sigmiod function operation;

(2) Replace the original BCE cross-entropy loss function with the VariFacalLoss loss function; in the process of replacing the BCE cross-entropy loss function with the VariFocalLoss loss function, the output layer of the model needs to be modified, and the VariFocalLoss loss function will introduce a learnable index γ, and the weight adjustment item in the calculation formula of the loss function has been modified to pay more attention to the learning of difficult samples. Therefore, the output layer needs to be modified accordingly to adapt to this change. In the improved YOLOX network model, the output The layer usually includes a classification branch and a regression branch. In the classification branch, each target needs to be classified, and in the regression branch, the location information of each target needs to be regressed. In order to adapt to the calculation of the VariFocalLoss loss function, each The prediction results of each target in the classification branch are processed. Specifically, the output of the classification branch needs to be processed by the sigmoid function first, and then it is converted into a prediction probability. According to this probability, the VariFocalLoss loss function is calculated. In the regression branch In , since the VariFocalLoss loss function only modifies the classification branch, the calculation method of the regression branch does not need to be changed;

(3) Training the improved YOLOX network model;

(4) Perform a pruning operation on the trained model, and evaluate the performance index of the model after pruning;

The attention mechanism module mines more available information through the space or channel dimension of the input feature for weighting processing, enhances the perception ability of the feature space and channel dimension, enables the network to have the ability to focus on inputting its features, and obtains better detection accuracy.

Using VariFocalLoss to replace the cross entropy in the original loss function can increase the weight of positive and negative samples and speed up the convergence of the model.

The step S6 includes the following steps;

First of all, it is necessary to evaluate and improve the performance of the YOLOX network model on unseen data. If the model cannot meet the performance requirements, adjust and improve, and further optimize the training process of the model by adjusting the training learning rate and the hyperparameters of the batch size. , after adjusting the model, it needs to be retrained and verified. This process requires multiple iterations until the model that meets the performance requirements is reached. Finally, if the improved pruned model can meet the performance requirements, the target is detected and analyzed. The target detection is It refers to the detection of the position and category of the target in the image or video. By deploying the improved pruned model, faster and more accurate target detection can be achieved, and the efficiency and accuracy in practical applications can be improved.

Through experimental verification, the method proposed by the present invention has achieved higher detection accuracy and speed in the target detection task. Compared with the traditional target detection method, the method proposed by the present invention can greatly improve the detection speed under the premise of maintaining high precision. Therefore, the present invention has higher practicality and economic benefit. Beneficial effects of the present invention:

The present invention provides a method based on YOLOX lightweight and network optimization, which inserts a channel-spatial attention mechanism module at the connection between feature extraction and enhancement layers to enhance feature extraction capabilities for targets of different scales and suppress redundancy information interference. At the same time, VariFocalLoss is used to replace the cross entropy in the original loss function, and weight values are added to positive and negative samples to control the shared weight of positive and negative samples on the total loss function value, so that the model can focus more on difficult samples during training. In order to solve the problem of unbalanced sample categories.

Although the optimization strategy in the present invention has effectively improved the object detection accuracy, further optimization is still needed to achieve end-to-end real-time object detection on mobile devices. In order to solve this problem, the present invention uses a pruning strategy to compress the model volume, reduce the calculation amount of the model, and realize end-to-end real-time target detection on the mobile device.

However, the pruning operation may affect the model accuracy, resulting in a decrease in detection accuracy. Therefore, the present invention verifies and analyzes the model after pruning. If the performance requirements can be met, the target is detected and analyzed; if the performance requirements cannot be met, the model is adjusted and improved until the performance requirements are met. In this way, it is possible to maintain a high target detection accuracy while ensuring the model is lightweight, and solve the problem of the evaluation performance degradation of the original YOLOX model after pruning.

Therefore, the present invention comprehensively applies multiple optimization strategies such as light weight, network optimization, loss function improvement, pruning, etc., effectively improves the end-to-end real-time target detection performance on mobile devices, and solves the problem of the original YOLOX in embedded devices. When deploying on the Internet, there are still problems such as large model size, high floating-point calculation, poor real-time performance, and decreased accuracy after pruning.

Description of drawings

Fig. 1 is a schematic flowchart of a method based on YOLOX lightweight and network optimization in the present invention.

Figure 2 is a schematic flow diagram of YOLOX network training.

Figure 3 is a schematic diagram of the model structure based on YOLOX network optimization.

Figure 4 is a flow chart based on YOLOX model training.

Figure 5 is a structural diagram of the CBAM attention mechanism.

Figure 6 is a schematic diagram of the accuracy of data set detection based on the improved YOLOX network model.

Detailed ways

The present invention will be described in further detail below in conjunction with the accompanying drawings.

In order to deepen the understanding of the present invention, the present invention will be further described below in conjunction with the accompanying drawings.

As shown in Figures 1-3: the present invention proposes a lightweight target detection method based on YOLOX, including the following steps:

S1 In the target detection task, prepare the data set required for training; the data set should include different types of target samples, and at the same time, the original image of the target in different scenes and different lighting conditions should be considered. Processing to ensure that the model can achieve accurate object detection in various environments;

The quality and quantity of the dataset are critical to the training and performance of the model. Therefore, it is necessary to select a dataset related to the object detection task, and preprocess and clean it for subsequent model training and evaluation.

S2 trains the original YOLOX neural network model on the data set, records and evaluates the performance indicators of the model;

Use the data set to train the original target detection model and record the performance indicators of the model. The performance indicators are accuracy rate, recall rate, and precision; where the accuracy rate indicates the proportion of positive samples that the model correctly predicts; the recall rate indicates that the model correctly detects Proportion of positive samples; accuracy means the proportion of all samples correctly predicted by the model; at the same time, adjust the learning rate and weight decay in the training process to improve the performance of the model; in the network training process, batch normalization technology is used , standardize the input data in each batch, accelerate the convergence speed of the model and improve the generalization ability of the model;

S3 performs a pruning operation on the original YOLOX neural network model; generates an improved YOLOX network model after pruning;

Reduce redundant parameters in the model through pruning technology, so as to achieve light weight; based on the preset pruning strategy, pruning operation is performed on the original network model to reduce the parameters and calculation amount in the model, so as to achieve light weight; pruning The final model needs to be retrained and its performance indicators should be evaluated to verify the impact of the pruning operation on the performance of the model. In the target detection task, pruning is used to remove some redundant network structures, thereby reducing the amount of calculation and improving the performance of the model. reasoning speed;

S4 trains and improves the YOLOX network model on the data set, improves the detection accuracy and evaluates the performance indicators of the model;

In order to improve the performance of the target detection model, a channel-spatial attention mechanism CBAM module is inserted between the backbone layer and the data enhancement layer of the original lightweight model, so that the model can automatically focus on important features related to target detection when processing images. Improve the attention and anti-interference ability of the network. Replace the original BCE cross-entropy loss function with the VariFocal Loss function to improve the classification performance of the model and the ability to identify difficult and easy samples. VariFocalLoss introduces variable key parameters in the loss function, making the model pay more attention to samples that are difficult to classify, thereby improving the classification performance of the model, training the improved model on the data set, and recording and evaluating the performance indicators of the model;

S5 performs a pruning operation on the improved YOLOX network model, and further adjusts the improved model to achieve higher detection accuracy and speed. In the target detection task, the lightweight and efficient deployment of the model is crucial, because they It directly affects the speed and accuracy of the model. The pruning operation is performed on the improved network model to further reduce the parameters and calculation amount of the model, so as to achieve faster and more efficient model deployment. At the same time, the pruned model can also be more suitable for scenarios with limited resources such as embedded devices , to improve the versatility and portability of the model;

S6 verifies and analyzes the improved pruned model to ensure that it can meet the performance requirements of the target detection task. If the requirements cannot be met, the improved model needs to be adjusted until the performance requirements are met.

Said S4 uses the process of improving the YOLOX network model, including the following steps:

S41: In order to enhance the feature extraction ability of the improved YOLOX network model for targets of different scales and suppress the interference of redundant information, a method of inserting a channel-spatial attention mechanism CBAM module between the backbone layer and the data enhancement layer channel is adopted. The CBAM module can adaptively learn the importance of channels and spaces in the feature map, and use this to weight and adjust the feature map, so that the network pays more attention to the feature information that contributes to target recognition and localization.

S42: In order to solve the problem of unbalanced sample categories, the method of replacing the original BCE cross-entropy loss function with the VariFacalLoss loss function is adopted. VariFocalLoss can add different weights to positive and negative samples, allowing the model to pay more attention to samples that are difficult to classify, so as to better balance sample categories during training and improve the classification accuracy of the model.

S43: In order to optimize the performance of the improved YOLOX network model, a means of training it is adopted. During the training process, the network continuously adjusts the weights and biases through the backpropagation algorithm to minimize the value of the loss function, thereby improving the classification and positioning capabilities of the model.

S44: In order to reduce the size of the model and the amount of calculation, so as to realize the end-to-end real-time object detection on the mobile device, the method of pruning the trained model is adopted. The pruning operation can remove unnecessary connections and parameters, thereby reducing the volume and computation of the model while maintaining the accuracy of the model. The pruned model needs to be retrained and its performance indicators evaluated to ensure that the model can still maintain high-precision target detection capabilities after pruning.

In said S43; the process of training the network model includes the following steps:

S431: collect the data set, and calibrate the original data set;

S432: divide the picture into a data set and a verification set according to a ratio of 8:2;

S433: Input the data set into the YOLOX neural network, input the predicted value and the actual value of the network into the loss function, obtain the loss value, and update the neural network parameters according to the gradient descent method;

S434: After each round of parameter update, the verification set is input into the YOLOX network for verification, and the loss degree of the verification machine is calculated.

S435: After every two iterations, input the pictures in the data set into the trained model to obtain the accuracy of the model;

S436: Repeat the above steps until the epoch reaches 300 rounds, and the model has converged.

The process of pruning the network model in the S3 includes the following steps:

S31: Determine the network layer to be pruned. Before model pruning, it is necessary to analyze and evaluate the network model to determine which network layers can be pruned. This process needs to consider factors such as the number of parameters of each network layer in the model, the amount of calculation, and the degree of contribution to the overall performance of the model. This can minimize the amount of calculation and volume of the model, while ensuring that the detection accuracy of the model will not be too low.

S32: sort the weights of the network layers to be pruned; the weight sorting of the network layers is one of the important steps of the pruning operation. Before pruning, it is necessary to sort all the parameters in the network to determine which parameters contribute less to the performance of the network and which parameters contribute more to the performance, and then select the network layer with less performance contribution for pruning.

S33: Determine the threshold according to the size of the weight. By sorting the weights of the network layers and setting thresholds, unnecessary or redundant parts of the network can be identified, and then pruned.

S34: Pruning the parts below the threshold to obtain new parameters. After determining the pruning threshold, the network parameters lower than the threshold are pruned. Specifically, the weights in the network are compressed, redundant and unnecessary parameters are eliminated, so as to achieve the purpose of reducing the amount of network storage and calculation.

S35: Save the new model parameters and weights, and generate a pruned model. After pruning, the number of model parameters is reduced, and the amount of calculation is also reduced accordingly, which can significantly improve the inference speed and efficiency of the model, and also help to achieve real-time object detection in resource-constrained scenarios such as embedded devices.

Below in conjunction with relevant background technology and implementation steps, the present invention will be further described:

In the step S1, select pictures using public datasets, including VOC2007 and VOC2012 datasets, including 21143 training images and corresponding xml files.

In the step S2, record performance evaluation indicators, including model accuracy mAP and model size Params(M), and use them as standards for subsequent performance evaluation.

YOLOX is the latest target detection algorithm of the YOLO series. It not only achieves the detection accuracy beyond the previous YOLO series, but also achieves a very competitive effect in end-to-end reasoning speed. However, when YOLOX is deployed on embedded devices, there are In order to solve the above problems and avoid unnecessary energy consumption caused by model pre-training, the lightweight method of YOLOX is proposed.

(1) Pruning the YOLOX network model

Read the network model trained by YOLOX, and save the model weight and structure. At this time, the model size is 3797KB. Determine the convolution layer to be cut. The determined layer is the C3_p4 layer, C3_n3 layer, and reduce_conv1 layer in the data enhancement layer And the bu_conv2 layer, sort the weights in each layer respectively, multiply the maximum value of the weight by the set pruning rate (40%) as the threshold of the pruning weight, and reset the weight of neurons below this threshold to 0 , weight neurons higher than the threshold, save new parameters and weights, and generate a new model structure after pruning;

Use the original model to directly perform lightweight pruning. Although the model can be compressed and the detection speed is improved, there is no difficulty in the process of downloading the embedded board, but it will directly cause a significant drop in the performance evaluation results of the model. Therefore, it is necessary to optimize the network model. Improve and improve, so the loss function is replaced while introducing the attention mechanism;

(1) Improve and optimize the YOLOX network model

In step S41, as shown in Fig. 5, the fused convolutional block attention mechanism CBAM is a combination of channel attention mechanism and spatial attention mechanism. First, through the channel attention mechanism, the input feature map is subjected to average pooling and maximum pooling operations respectively to realize the spatial information of the aggregated feature map, and the generated average pooling feature Favg and maximum pooling feature Fmax pass through the shared network layer. After the shared network is applied to each feature, the average pooled feature and the maximum pooled feature are element-wise summed, and the combined features are output through the Sigmiod activation function to the channel attention map Mc. Spatial attention performs average pooling and maximum pooling operations on the feature map along the channel axis, so that the feature map is compressed in the channel dimension, and the two feature maps are spliced in the channel dimension to generate an effective feature map, and then After a 7X7 convolutional layer. Finally, the final channel attention map Ms is obtained through the Sigmiod function operation.

As shown in Figure 4, it is the result of improving the YOLOX network structure. In the CSPNet layer of the cross-stage local network, the input part of the channel attention mechanism of CBAM is connected, and the data enhancement layer is connected in the output part of the spatial attention mechanism of CBAM.

The introduction of the attention mechanism inserts the channel-spatial attention mechanism module at the connection between the feature extraction and the enhancement layer, and screens out effective features from the channel and space dimensions, suppresses irrelevant features, enhances the expressive ability of features, and improves the The recognition accuracy of the model.

In step S43, in YOLOX object detection, the cross-entropy loss function has the problem of extreme imbalance between the object category and the background category. Using Focal loss can effectively solve the problem of imbalance between the target class and the background class. Focal loss formula is as follows:

Among them, p is the predicted probability of the target class, the range is [-1, 1]; y is the real positive and negative sample category, the value is 1 or -1; α is an adjustable scaling factor; (1-p) β times The square is the target modulation factor, and the β power of p is the background modulation factor. The two types of modulation factors can reduce the contribution of simple samples and increase the importance of false detection samples. Enables Focal loss to use a weighting method to solve the problem of category imbalance during training.

Focal loss handles positive and negative samples in an equal manner. In actual detection, the contribution of positive samples is more important. For this reason, Focal loss is further improved. Varifocal loss is based on cross-entropy binary, and the weighted method of Focal loss is used to deal with category imbalance in training. question. The binary formula for cross entropy is:

Among them, p is the predicted value, which represents the target score: q is the classification condition, and for the target class, set the positive sample category q value to the value between the pre-selection box and IoU, otherwise set it to 0. For the background category, the target q-value is 0 for all classes. As shown in the above formula, Varifocal loss uses the β power scaling factor of p to process negative samples, but not to scale positive samples. This can highlight the contribution of positive samples.

(3) Improved target detection network model

The present invention involves three parts: attention mechanism, loss function and model pruning. The attention mechanism is introduced between the backbone feature extraction layer and the data enhancement layer, so that the network has the ability to focus on inputting its features and obtain better detection accuracy. In the loss function part, the BCE cross-entropy loss function is replaced by the VariFocalLoss function, which increases the attention to difficult samples in the data set and achieves sample balance. The pruning operation of the improved target detection network model realizes the lightweight of the YOLOX network model, and solves the problems of large model size, high floating-point calculations, and poor real-time performance when the original YOLOX network is deployed on embedded devices The problem. (As shown in Figure 6)

To sum up, the present invention mainly solves the technical problems in two aspects. One is that the original YOLOX still has the problems of large model size, high amount of floating-point calculations and poor real-time performance when it is deployed on embedded devices; It is aimed at the problem of the evaluation performance decline after pruning of the original YOLOX model;

What is disclosed above is only an example of the present invention, which certainly cannot limit the scope of the present invention. Those of ordinary skill in the art can understand the process of realizing the above example, and make equivalent changes according to the claims of the present invention.

Claims (6)

1. A method based on the weight reduction of YOLOX and network optimization is characterized by comprising the following steps of;

s1: in the target detection task, preparing a data set required during training; the data set selects original images of the target under different scenes and different illumination conditions, and the model can realize accurate target detection under various environments through preprocessing the original images;

s2: training an original YOLOX neural network model on the data set, and recording and evaluating performance indexes of the model;

s3: performing pruning operation on the original YOLOX neural network model to generate a pruned improved YOLOX neural network model;

s4: training on the dataset to generate a pruned improved YOLOX network model;

s5: performing pruning operation on the improved YOLOX network model, and further adjusting to achieve higher detection accuracy and speed;

s6: verifying and analyzing the improved YOLOX network after improved pruning, and if the requirements on performance can be met, detecting and analyzing the target; if the performance requirements cannot be met, the improvement model is adjusted until the performance requirements are met.

2. The YOLOX-based lightweight and network optimization method according to claim 1, wherein in the step S1, a data set required for training is prepared, comprising the steps of:

(1) Collecting data: collecting a paspal VOC dataset from an open dataset officer network, including JPEGImages, imageSets and anotations; wherein JPEGImas contains a training dataset, imageSets contain each type of train. Txt, and val. Txt files, and anots contains xml files for each type;

(2) Data preprocessing: the data preprocessing is to preprocess the collected original image so as to adapt to model training; firstly, the original image is adjusted to be a specified size so as to facilitate subsequent processing; next, the color image is converted to a gray scale image, and finally, the pixel values in the image are scaled to between 0 and 1.

3. The YOLOX-based lightweight and network optimization method according to claim 1, wherein the step S2 trains an original YOLOX neural network model on a dataset, comprising the steps of:

(1) Adopting the data set preprocessed in the step S1;

(2) Dividing the data set into a training set and a verification set;

(3) Inputting the preprocessed data set into an original YOLOX neural network model, inputting a predicted value pred and a true value gt of the network into a loss function L, and obtaining a loss value through the following formula

Loss＝L(pred,gt)

Wherein, L represents a loss function, pred represents a predicted value of network output, gt represents a true value, optimizing network parameters according to the loss function L, updating the neural network parameters by using a gradient descent method, and setting the current neural network parameters as theta, wherein the updating formula is as follows:

wherein, eta represents the learning rate,

represents the gradient of the loss function L to the parameter θ, θ represents the parameter value of the t-th time step, θ _t+1 Parameter values representing the t+1st time step, updating neural network parameters through a plurality of iterations;

(4) After a round of parameter updating, the model is checked by using a verification set to verify the generalization capability of the model, the verification set is input into a YOLOX network, a loss measurement between a predicted result and a real result is calculated, the loss degree of the verification set is set as N, and a prediction frame of an ith sample is p _i The real frame is t _i The validation set loss L is calculated as follows:

where S is the number of prediction frames per grid, C is the number of target categories,

and->

Predicted and true values of the c-th class in the j-th lattice of the i-th sample,/-respectively>

And->

The j-th lattice of the i-th sample is respectively represented by whether the predicted value and the true value of the target exist or not,/->

And->

Confidence predictors and true values, pos, in the j-th grid of the i-th sample, respectively _ij Index set representing prediction frame with maximum cross-over ratio with real frame in j-th lattice of i-th sample

And->

Is two weight coefficients for balancing the weights of the lattices in which the object exists and the lattices in which the object does not exist;

(5) Each iteration is carried out twice, the pictures in the data set are input into an optimized YOLOX target detection network for training, and the accuracy of the model is obtained;

(6) Repeating the above steps until training is finished.

4. The YOLOX-based lightweight and network optimization method according to claim 1, wherein the step S3 performs pruning operation on the network trained in the step S2, and the method comprises the following steps:

(1) Determining a layer to be pruned according to the importance index of the network layer; for each network layer, calculating the sensitivity of one network layer to forward propagation on the model, and under the condition of given input, calculating the partial derivative of the output variable quantity to the weight of the layer, so as to obtain the sensitivity of the layer to the output, wherein the higher the sensitivity is, the higher the influence of the layer to the output is, and the priority should be given when pruning;

(2) Respectively carrying out weight sequencing on the network layers to be pruned;

(3) Determining a threshold according to the weight sequencing result and the pruning rate in each pruning layer;

(4) Rejecting weights lower than a threshold in a network, and reserving weights higher than the threshold;

(5) And (5) saving the new model parameters and weights, and generating the pruned improved YOLOX network model.

5. The YOLOX-based lightweight and network optimization method according to claim 1, wherein the step S4 improves the YOLOX network model, comprising the steps of:

(1) Inserting a channel-space attention mechanism (CBAM) module between a yolox trunk layer and a data enhancement layer channel;

firstly, respectively carrying out average pooling and maximum pooling operation on an input feature map through a channel attention mechanism to realize space information of an aggregate feature map and generate average pooling feature F _avg And maximize pooling feature F _max Through a shared network layer, after the shared network is applied to each feature, element summation is carried out on the average pooling feature and the maximum pooling feature, and the combined features are subjected to attention mapping Mc through a Sigmiod activation function output channel; carrying out average pooling and maximum pooling operation on the feature images along the channel axis by using the spatial attention, compressing the feature images in the channel dimension, splicing the two feature images in the channel dimension to generate an effective feature image, and then passing through a 7X7 convolution layer; finally, obtaining a final channel attention map Ms through Sigmiod function operation;

(2) Replacing the original BCE cross entropy loss function with a VarifacalLoss loss function; in replacing the BCE cross entropy loss function with the varifolloss function, the output layer of the model needs to be modified, the varifolloss function introduces a learnable index gamma, and the weight adjustment item in the calculation formula of the loss function is modified, so that more emphasis is placed on the learning of a difficult sample, in improving the YOLOX network model, the output layer generally comprises a classification branch and a regression branch, in the classification branch, each target needs to be classified, in the regression branch, the regression of position information needs to be carried out on each target, the output of the classification branch is firstly processed by the sigmoid function and then changed into a prediction probability, and the varifolloss function is calculated according to the probability;

(3) Training an improved YOLOX network model;

(4) And executing pruning operation on the trained model, and evaluating performance indexes of the model after pruning.

6. The YOLOX-based lightweight and network optimization method according to claim 1, wherein the step S6 comprises the steps of;

firstly, evaluating the performance of an improved YOLOX network model on unseen data, if the model cannot meet the performance requirement, adjusting and improving, further optimizing the training process of the model by adjusting the learning rate of training and the super parameters of batch size, and after the model is adjusted, retraining and verifying again, wherein the process needs to be iterated for a plurality of times until the model meeting the performance requirement is achieved, and finally, if the model after pruning is improved can meet the performance requirement, detecting and analyzing the target.

Images