A method based on YOLOX lightweight and network optimization

Abstract

The present invention discloses a method based on YOLOX lightweight and network optimization, comprising the following steps: S1: in a target detection task, preparing a data set required for training; S2: training an original YOLOX neural network model on the data set, recording and evaluating the performance indicators of the model; S3: performing a pruning operation on the original YOLOX neural network model to generate an improved YOLOX network model after pruning; S4: training and generating an improved YOLOX network model after pruning on the data set; S5: performing a pruning operation on the improved YOLOX network model; S6: verifying and analyzing the improved YOLOX network after pruning, and if the performance requirements are met, performing detection analysis on the target; if the performance requirements are not met, adjusting the improved model until the performance requirements are met. The present invention has high detection accuracy and speed in target detection, is easier to deploy and integrate in actual application scenarios, and also makes the reasoning process of the model more efficient and stable.

Description

Technical Field

[0001] The present invention belongs to the technical field of target detection in images, and specifically relates to a method based on YOLOX lightweighting and network optimization. Background Art

[0002] The object detection problem involves determining the location of objects in a given image and the category to which each object belongs (i.e., object localization and object classification). Object detection technology has found important applications in agriculture, healthcare, automated production, and other fields. Object detection primarily utilizes deep learning methods, which fall into two main categories. One category involves two-stage object detection algorithms based on candidate regions, including R-CNN, SPP-Net, Fast-RCNN, and Faster-RCNN.

[0003] Because these two-stage algorithms require the generation of a large number of candidate regions, object detection is relatively time-consuming. They first need to generate a large number of candidate regions and then classify and locate the objects within these candidate regions. This process consumes a large amount of computing resources and time, and is particularly slow in high-resolution images and complex scenes. Compared to two-stage algorithms, single-stage algorithms only require dense sampling in the image and then directly perform object detection at the sampled points. Therefore, they do not need to generate a large number of candidate regions, are faster, and are suitable for real-time object detection. Therefore, single-stage algorithms are more popular in practical applications.

[0004] YOLOX, an excellent single-target detection algorithm, can regress object category probabilities and location coordinates with high speed. However, direct use of YOLOX for object detection inevitably presents some challenges. The YOLOX model is large and inference speed is slow, so it needs to be lightweight. While achieving lightweight YOLOX networks, it is crucial to maintain their original accuracy. Therefore, network optimization is required after lightweighting to maintain detection accuracy. Summary of the Invention

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

[0006] In order to achieve the above object, the technical solution adopted by the present invention is:

[0007] A method based on YOLOX lightweighting and network optimization, comprising the following steps:

[0008] S1: In the target detection task, prepare the dataset required for training; the dataset selects original images in different scenes and different lighting conditions, and preprocesses the original images to ensure that the model can achieve accurate target detection in various environments;

[0009] S2: Train the original YOLOX neural network model on the dataset, and record and evaluate the performance metrics of the model;

[0010] S3: Perform pruning operations on the original YOLOX neural network model to generate an improved YOLOX network model after pruning;

[0011] S4: Train the pruned improved YOLOX network model on the dataset;

[0012] S5: Perform pruning operations on the improved YOLOX network model and make further adjustments to achieve higher detection accuracy and speed;

[0013] S6: Verify and analyze the improved YOLOX network after improved pruning. If it meets the performance requirements, perform detection analysis on the target; if it does not meet the performance requirements, adjust the improved model until it meets the performance requirements.

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

[0015] (1) Data collection: The PASCAL VOC dataset, including JPEGImages, ImageSets, and Annotations, was collected from the official website of the open dataset. JPEGImages contains the training dataset, ImageSets contains the train.txt, trainval.txt, and val.txt files of each type, and Annotations contains the XML files of each type.

[0016] (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. The original image size of the selected target detection is 416x416 pixels. Next, the color image is converted to a grayscale image to reduce the complexity of data storage and processing, and can reduce the time and computational complexity of model training. Finally, the pixel values in the image are scaled to between 0 and 1, making the data more stable during processing and reducing the gradient explosion and gradient disappearance problems during training.

[0017] Step S2 trains the original YOLOX neural network model on the dataset, including the following steps:

[0018] (1) Using the dataset preprocessed in step S1;

[0019] (2) The dataset is divided into training set and validation set in a ratio of 8:2;

[0020] (3) Input the preprocessed data set into the original YOLOX neural network model, and input the network's predicted value pred and the true value gt into the loss function L. The loss value is calculated using the following formula

[0021] Loss = L(pred,gt)

[0022] Where L represents the loss function, pred represents the predicted value of the network output, and gt represents the true value. The network parameters are optimized according to the loss function L, and the gradient descent method is used to update the neural network parameters. If the current neural network parameters are θ, the update formula is:

[0023]

[0024] Among them, η represents the learning rate, represents the gradient of the loss function L with respect to the parameter θ, where θ represents the parameter value at the t-th time step, θ t+1 Represents the parameter value of the t+1th time step. Through multiple iterations, the neural network parameters are updated to optimize the network performance and improve the accuracy and speed of target detection.

[0025] (4) After a round of parameter update, the model needs to be tested using the validation set to verify the generalization ability of the model. Specifically, the validation set is input into the YOLOX network and the loss metric between the predicted result and the true result is calculated, that is, the validation set loss. Assume that the validation set size is N and the prediction box of the i-th sample is p i , the real box is t i , then the validation set loss L can be calculated as follows:

[0026]

[0027] Among them, S is the number of prediction boxes for each grid, C is the number of target categories, and They represent the predicted value and true value of the cth category in the jth grid of the i-th sample, respectively. and They respectively represent the predicted value and true value of whether the target exists in the j-th grid of the i-th sample, and Represents the confidence prediction value and true value in the jth grid of the i-th sample, pos ij Represents the index set of the predicted box with the maximum intersection-over-union ratio with the true box in the jth grid of the i-th sample and are two weight coefficients used to balance the weights of grids with targets and grids without targets;

[0028] By calculating the validation set loss, you can evaluate the performance of the current model. If the loss is high, you need to continue training until the predetermined stopping condition is reached.

[0029] (5) After two iterations, the images in the dataset are input into the optimized YOLOX object detection network for training to obtain the accuracy of the model. The optimized YOLOX object detection network has higher detection accuracy and faster detection speed.

[0030] (6) Repeat the above steps until the training is completed.

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

[0032] (1) Determine the layers to be pruned based on the importance index of the network layer; for each network layer, calculate the sensitivity of the network layer to the forward propagation of the model. Under the given input, calculate the partial derivative of the change in the output with respect to the weight of the layer, thereby obtaining the sensitivity of the layer to the output. The greater the sensitivity, the greater the impact of the layer on the output, and the layer should be given priority when pruning;

[0033] (2) Sort the weights of the network layers to be pruned;

[0034] (3) Determine the threshold value based on the weight ranking results and pruning rate in each pruning layer;

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

[0036] (5) Save the new model parameters and weights to generate the pruned improved YOLOX network model.

[0037] The step S4 improves the YOLOX network model, including the following steps:

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

[0039] The CBAM module is an implementation of the channel-spatial attention mechanism, which can effectively improve the accuracy of the model. It mainly consists of two parts: channel attention and spatial attention.

[0040] First, the input feature map is subjected to average pooling and maximum pooling operations respectively through the channel attention mechanism to aggregate the spatial information of the feature map and generate the average pooling feature F avgand the maximum pooling feature F max After applying the shared network to each feature through the shared network layer, the average pooling feature and the maximum pooling feature are element-wise summed, and the combined feature is output as a channel attention map Mc through the Sigmiod activation function; spatial attention performs average pooling and maximum pooling operations on the feature map along the channel axis to compress the feature map in the channel dimension, and then splices the two feature maps in the channel dimension to generate a valid feature map, which is then passed through a 7X7 convolution layer; finally, the final channel attention map Ms is obtained through the Sigmiod function operation;

[0041] (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, it is necessary to modify the output layer of the model. The VariFocalLoss loss function introduces a learnable exponent γ and modifies the weight adjustment term in the calculation formula of the loss function, thereby paying 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 layer usually includes a classification branch and a regression branch. In the classification branch In the branch, each target needs to be classified, and in the regression branch, the position information of each target needs to be regressed. In order to adapt to the calculation of the VariFocalLoss loss function, the prediction results of each target in the classification branch need to be processed. Specifically, the output of the classification branch needs to be processed by the sigmoid function first, and then converted into a predicted probability. The VariFocalLoss loss function is calculated based on this probability. In the regression branch, since the VariFocalLoss loss function only modifies the classification branch, the calculation method of the regression branch does not need to be changed;

[0042] (3) Train the improved YOLOX network model;

[0043] (4) Perform pruning operations on the trained model and evaluate the performance indicators of the pruned model;

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

[0045] Using VariFocalLoss instead of the cross entropy in the original loss function can increase the weights of positive and negative samples and speed up the convergence of the model.

[0046] The step S6 comprises the following steps:

[0047] First, it is necessary to evaluate the performance of the improved YOLOX network model on unseen data. If the model cannot meet the performance requirements, adjustments and improvements are made. By adjusting the training learning rate and batch size hyperparameters, the model training process is further optimized. After adjusting the model, it is necessary to retrain and verify. This process requires multiple iterations until a model that meets the performance requirements is achieved. Finally, if the improved pruned model can meet the performance requirements, target detection analysis is performed. Target detection refers to detecting the location and category of the target in an image or video. By deploying the improved pruned model, faster and more accurate target detection can be achieved, thereby improving efficiency and accuracy in practical applications.

[0048] Experimental verification shows that the method proposed in this invention achieves high detection accuracy and speed in target detection tasks. Compared with traditional target detection methods, the method proposed in this invention can significantly improve detection speed while maintaining high accuracy. Therefore, this invention has high practicality and economic benefits. Beneficial effects of the invention:

[0049] This paper provides a YOLOX-based lightweight and network optimization method. This method inserts a channel-spatial attention mechanism module at the junction of the feature extraction and enhancement layers to enhance the feature extraction capability of objects of different scales and suppress the interference of redundant information. At the same time, VariFocalLoss is used to replace the cross entropy in the original loss function, and weights are added to positive and negative samples to control the shared weight of the total loss function value. This allows the model to focus more on difficult-to-distinguish samples during training, thereby addressing the problem of sample class imbalance.

[0050] Although the optimization strategy in this paper has effectively improved object detection accuracy, further optimization is still needed to achieve end-to-end real-time object detection on mobile devices. To address this problem, this paper uses a pruning strategy to compress the model size and reduce the model computational complexity, achieving end-to-end real-time object detection on mobile devices.

[0051] However, pruning can affect model accuracy, reducing detection precision. Therefore, the present invention verifies and analyzes the model after pruning. If it meets performance requirements, it then performs target detection analysis. If not, it adjusts and improves the model until it does. This ensures a lightweight model while maintaining high target detection accuracy, resolving the issue of decreased evaluation performance in the original YOLOX model after pruning.

[0052] Therefore, the present invention comprehensively applies multiple optimization strategies such as lightweighting, network optimization, loss function improvement, and pruning, effectively improving the end-to-end real-time target detection performance on mobile devices, and solving the problems of large model size, high floating-point calculation volume, poor real-time performance, and decreased accuracy after pruning when the original YOLOX is deployed on embedded devices. BRIEF DESCRIPTION OF THE DRAWINGS

[0053] Figure 1 It is a flow chart of a method of lightweighting and network optimization based on YOLOX in the present invention.

[0054] Figure 2 It is a flowchart of YOLOX network training.

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

[0056] Figure 4 It is a flowchart based on YOLOX model training.

[0057] Figure 5 This is the structural diagram of the CBAM attention mechanism.

[0058] Figure 6 This is a schematic diagram of the accuracy of data set detection based on the improved YOLOX network model. DETAILED DESCRIPTION

[0059] The present invention will be described in further detail below with reference to the accompanying drawings.

[0060] In order to deepen the understanding of the present invention, the present invention will be further described below with reference to the accompanying drawings.

[0061] like Figure 1-Figure 3 As shown: The present invention proposes a target detection method based on YOLOX lightweight, comprising the following steps:

[0062] S1: In the target detection task, prepare the dataset required for training. The dataset should include target samples of different categories, and consider the original images of the target in different scenes and different lighting conditions. By preprocessing the original images, ensure that the model can achieve accurate target detection in various environments.

[0063] The quality and quantity of the dataset are crucial 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 to facilitate subsequent model training and evaluation.

[0064] S2 trains the original YOLOX neural network model on the dataset, and records and evaluates the performance metrics of the model;

[0065] The original object detection model was trained using the dataset, and the model's performance indicators, including accuracy, recall, and precision, were recorded. Accuracy represents the proportion of positive samples correctly predicted by the model; recall represents the proportion of positive samples correctly detected by the model; and precision represents the proportion of all samples correctly predicted by the model. The learning rate and weight decay were adjusted during training to improve the model's performance. During network training, batch normalization was used to standardize the input data in each batch, accelerating the model's convergence and improving its generalization capabilities.

[0066] S3 performs pruning operations on the original YOLOX neural network model to generate an improved YOLOX network model after pruning;

[0067] Pruning technology is used to reduce redundant parameters in the model, thereby achieving lightweighting. Based on a pre-set pruning strategy, the original network model is pruned to reduce the parameters and computational complexity in the model, thereby achieving lightweighting. The pruned model needs to be retrained and its performance indicators evaluated to verify the impact of the pruning operation on model performance. In object detection tasks, pruning is used to remove some redundant network structures, thereby reducing computational complexity and improving the model's inference speed.

[0068] S4 trains and improves the YOLOX network model on the dataset to improve the detection accuracy and performance indicators of the evaluation model;

[0069] To improve the performance of the target detection model, a channel-spatial attention mechanism (CBAM) module was inserted between the backbone layer and the data enhancement layer of the original lightweight model. This enabled the model to automatically focus on important features related to target detection when processing images, improving the network's attention and anti-interference capabilities. The original BCE cross-entropy loss function was replaced with the VariFocal Loss loss function to improve the model's classification performance and its ability to identify difficult and easy samples. VariFocalLoss introduces a variable focus parameter into the loss function, allowing the model to pay more attention to difficult-to-classify samples, thereby improving the model's classification performance. The improved model was trained on the dataset, and the model's performance indicators were recorded and evaluated.

[0070] S5 performs pruning on the improved YOLOX network model and further adjusts the improved model to achieve higher detection accuracy and speed. In object detection tasks, lightweight and efficient model deployment are crucial because they directly affect the speed and accuracy of the model. Pruning the improved network model further reduces the model's parameters and computational complexity, thereby achieving faster and more efficient model deployment. At the same time, the pruned model can also be more suitable for resource-limited scenarios such as embedded devices, improving the model's versatility and portability.

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

[0072] The process of improving the YOLOX network model used by S4 includes the following steps:

[0073] S41: To enhance the improved YOLOX network model's ability to extract features from objects of varying scales and suppress the interference of redundant information, a channel-spatial attention mechanism (CBAM) module is inserted between the backbone layer and the data augmentation layer. The CBAM module adaptively learns the importance of channels and spaces in feature maps and uses this to weight and adjust feature maps, allowing the network to focus more on feature information that contributes to object recognition and localization.

[0074] S42: To address the imbalance in sample categories, we replaced the original BCE cross entropy loss function with the VariFocalLoss loss function. VariFocalLoss assigns different weights to positive and negative samples, allowing the model to focus more on difficult-to-classify samples. This allows for better balancing of sample categories during training and improves the model's classification accuracy.

[0075] S43: To optimize the performance of the improved YOLOX network model, training is performed. During the training process, the network continuously adjusts weights and biases through the backpropagation algorithm to minimize the value of the loss function, thereby improving the model's classification and positioning capabilities.

[0076] S44: To reduce model size and computational complexity, and to facilitate end-to-end real-time object detection on mobile devices, a trained model is pruned. Pruning removes unnecessary connections and parameters, reducing model size and computational complexity while maintaining model accuracy. The pruned model is retrained and its performance metrics evaluated to ensure that it maintains high-precision object detection capabilities after pruning.

[0077] In S43, the process of training the network model includes the following steps:

[0078] S431: Collecting a data set and calibrating the original data set;

[0079] S432: Divide the images into a data set and a validation set in a ratio of 8:2;

[0080] S433: Input the data set into the YOLOX neural network, input the network's predicted value and true value into the loss function, calculate the loss value, and update the neural network parameters according to the gradient descent method;

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

[0082] S435: After each two iterations, the images in the dataset are input into the trained model to obtain the accuracy of the model;

[0083] S436: Repeat the above steps until the epoch reaches 300 rounds, at which point the model has converged.

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

[0085] S31: Determine the network layers to prune. Before pruning the model, the network model needs to be analyzed and evaluated to determine which network layers can be pruned. This process considers factors such as the number of parameters in each network layer, the computational effort, and the contribution to the overall model performance. This minimizes the computational effort and size of the model while ensuring that the model's detection accuracy does not decrease significantly.

[0086] S32: Sort the weights of the network layers to be pruned. Sorting the weights of network layers is an important step in pruning. Before pruning, all parameters in the network need to be sorted to determine which parameters contribute less to the network's performance and which parameters contribute more to the performance. Then, the network layers with the lesser contribution are selected for pruning.

[0087] S33: Determine the threshold value based on the weight. By sorting the weights of the network layers and setting the threshold value, it is possible to identify unnecessary or redundant parts of the network and then perform pruning.

[0088] S34: Prune the network parameters below the threshold to obtain new parameters. After determining the pruning threshold, prune the network parameters below the threshold. Specifically, the network weights are compressed to remove redundant and unnecessary parameters, thereby reducing network storage and computational complexity.

[0089] S35: Save the new model parameters and weights to generate a pruned model. The pruned model has fewer parameters and a correspondingly smaller computational effort, significantly improving the model's inference speed and efficiency. This also helps achieve real-time object detection in resource-constrained scenarios such as embedded devices.

[0090] The present invention is further described below in conjunction with relevant background technologies and implementation steps:

[0091] In step S1, images from public datasets are selected, including VOC2007 and VOC2012 datasets, which contain 21,143 training images and corresponding xml files.

[0092] In step S2, the performance evaluation indicators, including model accuracy mAP and model size Params (M), are recorded and used as the criteria for subsequent performance evaluation.

[0093] YOLOX is the latest object detection algorithm in the YOLO series. It not only achieves detection accuracy that surpasses previous YOLO series algorithms, but also achieves highly competitive end-to-end inference speed. However, when YOLOX is deployed on embedded devices, it faces problems such as large model size, high floating-point computing workload, and poor real-time performance. To address these issues while avoiding unnecessary energy consumption caused by model pre-training, a lightweight approach to YOLOX has been proposed.

[0094] (1) Pruning the YOLOX network model

[0095] The network model trained by YOLOX is read and the model weights and structure are saved. The model size is 3797KB at this time. The convolutional layers to be pruned are determined. The layers determined are the C3_p4 layer, C3_n3 layer, reduce_conv1 layer, and bu_conv2 layer in the data enhancement layer. The weights in each layer are sorted respectively. The maximum weight multiplied by the set pruning rate (40%) is used as the threshold for pruning weights. The weights of neurons below the threshold are reset to 0, and the weight neurons above the threshold are saved. The new parameters and weights are generated to generate a new model structure after pruning.

[0096] Using the original model for lightweight pruning can compress the model, improve the detection speed, and facilitate downloading to embedded boards. However, this can lead to a significant drop in model performance evaluation results. Therefore, the network model needs to be refined and improved. Therefore, the attention mechanism is introduced and the loss function is replaced.

[0097] (1) Improve and optimize the YOLOX network model

[0098] In step S41, Figure 5As shown in the figure, the fused convolutional block attention mechanism (CBAM) combines the channel attention mechanism and the spatial attention mechanism. First, the channel attention mechanism performs average pooling and maximum pooling operations on the input feature map to aggregate the spatial information of the feature map. The generated average pooling features Favg and maximum pooling features Fmax are passed through a shared network layer. After the shared network is applied to each feature, the average pooling features and maximum pooling features are element-wise summed. The merged features are then activated using the Sigmiod function to output the channel attention map Mc. Spatial attention performs average pooling and maximum pooling operations on the feature map along the channel axis, compressing the feature map in the channel dimension. The two feature maps are then concatenated in the channel dimension to generate a valid feature map, which is then passed through a 7x7 convolutional layer. Finally, the final channel attention map Ms is obtained through the Sigmiod function.

[0099] like Figure 4 The figure shows the improved results of the YOLOX network structure. In the cross-stage local network CSPNet layer, the channel attention mechanism input of CBAM is connected, and the data enhancement layer is connected to the spatial attention mechanism output of CBAM.

[0100] The introduction of the attention mechanism inserts a channel-space attention mechanism module at the connection between the feature extraction and enhancement layers, which filters out effective features from the channel and spatial dimensions respectively, suppresses irrelevant features, enhances the expressiveness of features, and improves the recognition accuracy of the model.

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

[0102]

[0103] Where p is the predicted probability of the target class, ranging from -1 to 1; y is the true positive and negative sample category, taking the value of 1 or -1; α is an adjustable scaling factor; (1-p) to the power of β is the target class modulation factor, and p to the power of β is the background class modulation factor. These two modulation factors can reduce the contribution of simple samples and increase the importance of false positive samples. This allows Focal Loss to use a weighted approach to address class imbalance during training.

[0104] Focal loss treats positive and negative samples equally, but in actual detection, the contribution of positive samples is more important. Therefore, Focal loss is further improved. Varifocal loss is based on cross-entropy binary and borrows the weighted method of Focal loss to deal with the class imbalance problem in training. The cross-entropy binary formula is:

[0105]

[0106] Where p is the predicted value, representing the objectness score, and q is the classification condition. For the target class, the q-value of the positive class is set to the value between the preselected bounding box and the IoU; otherwise, it is set to 0. For the background class, the target q-value of all classes is 0. As shown in the above formula, Varifocal Loss uses a β-power scaling factor of p to process negative samples, while not scaling positive samples. This effectively emphasizes the contribution of positive samples.

[0107] (3) Improved target detection network model

[0108] 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, thereby obtaining 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 improved target detection network model is pruned to achieve the lightweight YOLOX network model, solving the problems of large model size, high floating-point calculation amount and poor real-time performance when the original YOLOX network is deployed on embedded devices. (such as Figure 6 shown)

[0109] In summary, the present invention mainly solves two technical problems: first, the problem that the original YOLOX model still has large model size, high floating-point calculation workload, and poor real-time performance when deployed on embedded devices; second, the problem that the evaluation performance of the original YOLOX model decreases after pruning;

[0110] The above disclosure is only an example of the present invention and certainly cannot be used to limit the scope of the present invention. Ordinary technicians in this field can understand the process of implementing the above example and make equivalent changes in accordance with the claims of the present invention.

Claims

1. A method based on YOLOX lightweight and network optimization, characterized in that: The following steps are included: S1: In the target detection task, prepare the dataset required for training; the dataset selects original images of the target in different scenes and different lighting conditions, and preprocesses the original images to ensure that the model can achieve accurate target detection in various environments; S2: Train the original YOLOX neural network model on the dataset, and record and evaluate the performance metrics of the model; S3: Perform pruning operations on the original YOLOX neural network model to generate an improved YOLOX network model after pruning; S4: Train the pruned improved YOLOX network model on the dataset; S5: Perform pruning operations 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 it meets the performance requirements, perform detection analysis on the target. If it does not meet the performance requirements, adjust the improved model until it meets the performance requirements. The step S4 improves the YOLOX network model, including the following steps: (1) Insert the channel-spatial attention mechanism CBAM module between the yolox backbone layer and the data enhancement layer channel; First, the input feature map is subjected to average pooling and maximum pooling operations respectively through the channel attention mechanism to aggregate the spatial information of the feature map and generate the average pooling feature F avg and the maximum pooling feature F max After applying the shared network to each feature through the shared network layer, the average pooling feature and the maximum pooling feature are element-wise summed, and the combined feature is output as the channel attention map Mc through the Sigmiod activation function; spatial attention performs average pooling and maximum pooling operations on the feature map along the channel axis to compress the feature map in the channel dimension, and then splices the two feature maps in the channel dimension to generate a valid feature map, which is then passed through a 7X7 convolution layer; finally, the final channel attention map Ms is obtained 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, it is necessary to modify the output layer of the model. The VariFocalLoss loss function introduces a learnable exponent γ and modifies the weight adjustment term in the calculation formula of the loss function, thereby paying more attention to the learning of difficult samples. In the improved YOLOX network model, the output 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 position information of each target needs to be regressed. The output of the classification branch is first processed by the sigmoid function and then converted into a predicted probability. The VariFocalLoss loss function is calculated based on this probability. (3) Train the improved YOLOX network model; (4) Perform pruning operations on the trained model and evaluate the performance indicators of the pruned model.

2. The method of lightweighting and network optimization based on YOLOX according to claim 1, characterized in that: In step S1, preparing the data set required for training includes the following steps: (1) Data collection: The PASCAL VOC dataset, including JPEGImages, ImageSets, and Annotations, was collected from the official website of the open dataset. JPEGImages contains the training dataset, ImageSets contains the train.txt, trainval.txt, and val.txt files of each type, and Annotations contains the XML files of each type. (2) Data preprocessing: Data preprocessing is to preprocess the collected original images to make them suitable for model training; first, the original images are resized to a specified size for subsequent processing; next, the color images are converted to grayscale images, and finally, the pixel values in the images are scaled to between 0 and 1.

3. The method of lightweighting and network optimization based on YOLOX according to claim 1, characterized in that: Step S2 trains the original YOLOX neural network model on the dataset, including the following steps: (1) Using the dataset preprocessed in step S1; (2) Divide the dataset into training set and validation set; (3) Input the preprocessed data set into the original YOLOX neural network model, and input the network's predicted value pred and the true value gt into the loss function L. The loss value is calculated using the following formula Loss = L(pred,gt) Where L represents the loss function, pred represents the predicted value of the network output, and gt represents the true value. The network parameters are optimized according to the loss function L, and the gradient descent method is used to update the neural network parameters. If the current neural network parameters are θ, the update formula is: Among them, η represents the learning rate, represents the gradient of the loss function L with respect to the parameter θ, where θ represents the parameter value at the t-th time step, θ t+1 Represents the parameter value at the t+1th time step, and updates the neural network parameters through multiple iterations; (4) After a round of parameter update, the model is tested using the validation set to verify the generalization ability of the model. The validation set is input into the YOLOX network to calculate the loss metric between the predicted result and the true result, the validation set loss degree, and the validation set size is set to N, and the prediction box of the i-th sample is p i , the real box is t i , then the validation set loss L is calculated as follows: Among them, S is the number of prediction boxes for each grid, C is the number of target categories, and They represent the predicted value and true value of the cth category in the jth grid of the i-th sample, respectively. and They respectively represent the predicted value and true value of whether the target exists in the j-th grid of the i-th sample, and Represents the confidence prediction value and true value in the jth grid of the i-th sample, pos ij Represents the index set of the predicted box with the maximum intersection-over-union ratio with the true box in the jth grid of the i-th sample and are two weight coefficients used to balance the weights of grids with targets and grids without targets; (5) After two iterations, the images in the dataset are input into the optimized YOLOX object detection network for training to obtain the accuracy of the model; (6) Repeat the above steps until the training is completed.

4. The method of lightweighting and network optimization based on YOLOX according to claim 1, characterized in that: In step S3, a pruning operation is performed on the network trained in step S2, including the following steps: (1) Determine the layers to be pruned based on the importance index of the network layer; for each network layer, calculate the sensitivity of the network layer to the forward propagation of the model. Under the given input, calculate the partial derivative of the change in the output with respect to the weight of the layer, thereby obtaining the sensitivity of the layer to the output. The greater the sensitivity, the greater the impact of the layer on the output, and the layer should be given priority when pruning; (2) Sort the weights of the network layers to be pruned; (3) Determine the threshold value based on the weight ranking results and pruning rate in each pruning layer; (4) Eliminate the weights below the threshold in the network and retain the weights above the threshold; (5) Save the new model parameters and weights to generate the pruned improved YOLOX network model.

5. The method of lightweighting and network optimization based on YOLOX according to claim 1, characterized in that: The step S6 includes the following steps: First, it is necessary to evaluate the performance of the improved YOLOX network model on unseen data. If the model cannot meet the performance requirements, adjustments and improvements are made. By adjusting the training learning rate and batch size hyperparameters, the model training process is further optimized. After adjusting the model, it is necessary to retrain and verify. This process requires multiple iterations until a model that meets the performance requirements is achieved. Finally, if the improved pruned model can meet the performance requirements, the target is detected and analyzed.

Images