A lightweight model compression vehicle detection method

Abstract

The application discloses a kind of light weight model compression vehicle detection methods, it is related to convolutional neural network computer vision field.The method uses channel pruning algorithm and model compression algorithm is combined, and YOLOv4-tiny light weight model is compressed to model.By using Mosaic data enhancement method to vehicle training data set is carried out data enhancement and expansion, and the original data set is expanded to the twice of original training data set;Using BN layer parameter to measure the importance of channel, channel pruning is carried out, and the channel below the pruning proportion in input layer weight is deleted, the corresponding channel of output layer is adjusted, and the parameter quantity of network model structure is reduced;The knowledge distillation algorithm based on classification task is applied to target detection, and local sample points are sampled and extracted, to learn local area information from teacher model, improve the accuracy of network model.

Description

TECHNICAL FIELD

[0001] The present application relates to the field of convolutional neural network computer vision, in particular to a lightweight model compression vehicle detection method. BACKGROUND

[0002] With the development of computer science and technology, the improvement of computing power and the wide application of computer vision principles in various scenes, using convolutional neural network to process target detection task has become the focus of research;

[0003] Traditional research mainly focuses on improving the detection accuracy of network model, and the performance of neural network model is constantly improved, while the complexity of network model is also increasing; the high-precision convolutional neural network model requires high computing power and memory consumption, which makes it almost impossible to use in small devices; for example: the two-stage target detection algorithm Faster-RCNN based on region extraction and the one-stage target detection algorithm YOLOv4 based on prior box classification cannot be deployed in resource-limited environments; although the lightweight model has lower deployment cost, it is still very difficult to deploy in small devices with very low computing power such as mobile devices, home computers and Raspberry Pi; for example: the Fps of YOLOv4-tiny model deployed on Raspberry Pi 4B is only 1. SUMMARY

[0004] In view of the deficiencies of the prior art, the present application provides a lightweight model compression vehicle detection method, which combines channel pruning algorithm and knowledge distillation algorithm to compress the YOLOv4-tiny lightweight model, improve the detection accuracy and speed of the algorithm; use BN layer parameters to measure the importance of channels, perform channel pruning and reduce the number of network model structure parameters; apply the knowledge distillation algorithm based on classification task to target detection, sample and extract local sample points to improve the accuracy of the network model;

[0005] The present application provides the following technical solutions:

[0006] A lightweight model compression vehicle detection method, the method comprising the following steps:

[0007] Step 1: collect the vehicle data set required for model training, and divide the data set into vehicle training data set and vehicle test data set according to the proportion;

[0008] Step 2: use data enhancement algorithm to perform data enhancement and expansion on the vehicle training data set, and expand the vehicle training data set to twice the original vehicle training data set;

[0009] Step 3: Put the vehicle training data set after data enhancement in step 2 into the YOLOv4-tiny model for training. After training converges, save the weight file after training is completed;

[0010] Step 4: Set the pruning ratio, select a specific value in the channel as the weight of each layer channel, judge its contribution to the model, delete the channels in the input layer weight below the pruning ratio, and adjust the output layer channels corresponding to the deleted input layer;

[0011] Step 5: Re-train the pruned YOLOv4-tiny model to get the weight file trained to convergence, and load it into the model for testing;

[0012] Step 6: Use the YOLOv4 model as the teacher model of the knowledge distillation algorithm, put the vehicle training data set after data enhancement in step 2 into the YOLOv4 model for training. After training converges, save the weight file after training is completed;

[0013] Step 7: The pruned YOLOv4-tiny model is used as the student model of the knowledge distillation algorithm. Put the vehicle training data into the teacher model and the student model respectively to get the feature map. Calculate the IoU value of the prior box and the real box, set the threshold, and calculate the cross entropy function of the prior box with IoU value greater than the threshold;

[0014] Step 8: Add the cross entropy function in step 7 and the loss function of YOLOv4-tiny, train the YOLOv4-tiny model, and save the weight file after training converges;

[0015] Step 9: Test the YOLOv4-tiny model completed by knowledge distillation on the vehicle test data set obtained in step 1 to determine the position coordinates and mark the corresponding positions in the original vehicle test data set;

[0016] The specific process of step 1: Use the vehicle video recorded by the vehicle data recorder during driving to convert the sequence in the video into a frame-by-frame picture as a data set. Label the detected objects in the picture according to the proportion as the vehicle training set and the vehicle test set;

[0017] The specific process of step 2: Use the Mosaic data enhancement method to mix four pictures through random cropping, random scaling and random arrangement to generate one picture for model training;

[0018] The specific value in step 4: the weight value of the BN (Batch Normalization) layer;

[0019] The method for learning feature information from the teacher model in step 7 is as follows: after the color picture passes through two model backbone networks, two feature maps of equal size are generated, but the channel numbers of the feature maps are not equal; adaptive feature map convolution is used to make the channel numbers of the student model and the teacher model equal through a convolution layer, and a loss function between the two feature maps is calculated, and the specific process is as follows:

[0020] 1) Introduce the YOLOv4 model trained in step 6 as the teacher model of the knowledge distillation algorithm;

[0021] 2) Introduce the YOLOv4-tiny model pruned in step 5 as the student model of the knowledge distillation algorithm;

[0022] 3) After the vehicle training data set enters the two models, two feature maps of the same size are obtained, and adaptive feature map convolution is used, that is, the channel numbers of the student model and the teacher model are made equal through a convolution layer;

[0023] 4) Calculate the Loss value of the two feature maps to make the student model learn the feature information of the teacher model;

[0024] The loss function calculation is shown in equation (1):

[0025]

[0026] Where N p represents the number of 1s in the generated matrix I, that is, the number of prior boxes near the position of the real box, W is the width of the feature map, H is the height of the feature map, C is the number of different size prior boxes, s ijc represents the value of the position of the feature map generated by the student model, t ijc represents the value of the position of the feature map generated by the teacher model.

[0027] The IoU calculation process in step 7 is shown in equation (2):

[0028]

[0029] Where S Truth is the area of the real box, and S box is the area of the prior box.

[0030] Beneficial technical effects:

[0031] 1. It can enrich the background and small targets of the detected objects in the data set. In the training process, four pictures are calculated at a time, making full use of the GPU in the device.

[0032] 2. The channel number of the network model is reduced by channel pruning, the parameter quantity of the network model structure is reduced, the detection speed of the lightweight model is improved, and the lightweight model can be deployed in a resource-limited environment;

[0033] 3. If the detection speed of the network model is improved only by model pruning, the detection accuracy of the network model will decrease; by using the knowledge distillation method, the local area near the real label frame is focused on, the information of the local area is used to learn knowledge from the teacher model, the Loss function is modified, and the knowledge transferred from the teacher model is more accurate, so that the detection accuracy is improved without changing the parameter quantity of the network model. BRIEF DESCRIPTION OF DRAWINGS

[0034] Figure 1 is a flowchart of a lightweight model compression vehicle detection method according to an embodiment of the application;

[0035] Figure 2 is a YOLOv4-tiny network model structure diagram of a lightweight model compression vehicle detection method according to an embodiment of the application;

[0036] Figure 3 is a YOLOv4-tiny network model structure diagram after channel pruning of a lightweight model compression vehicle detection method according to an embodiment of the application;

[0037] Figure 4 is a knowledge distillation algorithm diagram of a lightweight model compression vehicle detection method according to an embodiment of the application;

[0038] Among them, the upper part of the model is a complex teacher model of the network model which has been trained, and the lower part of the model is a simple student model of the structure which has not been trained; after the color picture passes through the two model backbone networks, two feature maps of equal size are generated, but the channel numbers of the feature maps are not equal; the adaptive feature map convolution in the figure is to make the channel number of the student model equal to the channel number of the teacher model by the convolution layer, and the loss function of the two feature maps is calculated;

[0039] Figure 5 is a local feature extraction diagram of a lightweight model compression vehicle detection method according to an embodiment of the application;

[0040] Figure 6 is an original picture in a test set of a lightweight model compression vehicle detection method according to an embodiment of the application;

[0041] Figure 7 is a detection result diagram of a lightweight model compression vehicle detection method according to an embodiment of the application. DETAILED DESCRIPTION

[0042] In order to make the purpose, technical scheme and advantages of the present application clearer, the present application will be further described in detail below with the aid of the drawings and specific embodiments; the specific embodiments described herein are only used to explain the present application and not used to limit the present application;

[0043] The present application first acquires vehicle video data recorded by a vehicle in the driving process, converts the video sequence into pictures one by one, labels the detected objects in the pictures, and divides the data set into training data set and test data set according to the proportion of 9:1; then the Mosaic data enhancement method is used to enhance and expand the training data set, and the original data set is expanded to twice the original training data set; the enhanced training data set is put into the YOLOv4-tiny model for training, and after the training converges, the weight file after the training is completed is saved; the trained model is pruned, the pruning ratio is defined, the specific value in the channel is selected as the weight of each layer channel, the contribution degree of the model is judged, and the channel with small weight is deleted; after a certain number of channels are deleted, the number of channels of the output layer will change, and the corresponding channels of the output layer need to be adjusted, the pruned YOLOv4-tiny model is retrained with a small learning rate, and the weight file trained to convergence is loaded into the model for testing;

[0044] Next, the pruned model is distilled, the YOLOv4 model is selected as the teacher model for knowledge distillation, the data enhanced model is put into the YOLOv4 model for training, and the training is converged; the training data set is put into the teacher model and the student model respectively to obtain the feature map; then the IoU value of the prior box and the real box is calculated, the threshold is set to 0.7 times the maximum value of the IoU value of all prior boxes and real boxes, the prior box with an IoU value greater than the threshold is calculated by the cross entropy function, so that the student model can learn the features of the teacher model more fully; the trained YOLOv4-tiny model is used to classify the test data set to determine the classification and position coordinates of the target object, and the corresponding labels are marked in the test data set;

[0045] Figure 1 is a specific flowchart of the lightweight model compression vehicle detection method of the present embodiment, which comprises the following steps:

[0046] Step 1: collect the vehicle data set required for model training, and divide the vehicle data set into vehicle training data set and vehicle test data set according to the proportion;

[0047] In the embodiment, the video of the vehicle and the pedestrian recorded by the vehicle during driving is used, the sequence in the video is changed into a picture by picture, the detected object in the picture is labeled, the data set has about 8000 pictures, the pictures in the data set are divided into two parts according to the ratio of 9:1, the part with more pictures is used as the training data set of the model, and the other part is used as the test data set of the model;

[0048] Step 2: Data augmentation is performed on the vehicle training data set by using data augmentation, and the vehicle training data set is expanded to twice the original training data set;

[0049] In the embodiment, the Mosaic data augmentation method is used, four pictures are mixed by random cropping, random scaling and random arrangement, and a picture is generated after mixing for model training; the specific process is as follows:

[0050] 1) four pictures are randomly extracted from the vehicle training data set, and a coordinate point is randomly generated in the picture size range;

[0051] 2) the generated coordinate point is cross divided by using the horizontal coordinate and the vertical coordinate, so that one picture is generated four parts with different areas;

[0052] 3) the four pictures are randomly cropped to the areas corresponding to the four parts generated in 2);

[0053] 4) the cropped pictures are randomly transformed in brightness, contrast and saturation, and then the four pictures are fused and placed in the positions corresponding to the cropped areas, and the four pictures are spliced into a new picture;

[0054] 5) the new picture is saved, and the classification and position of the detected object in the new picture are generated;

[0055] 6) repeat 1) to 5), and expand the training data set pictures to twice the original;

[0056] Step 3: the vehicle training data set after data augmentation in step 2 is put into the YOLOv4-tiny model for training, and after the training converges, the weight file after the training is completed is saved;

[0057] Since the YOLOv4-tiny model needs to be pruned, the vehicle training data set is sent into the model for training to generate the weight of each layer of the model, so as to judge the contribution degree of each channel of the model to the whole model, and after the training converges, the weight file is saved; the structure of the YOLOv4-tiny model is as shown in Figure 2 ;

[0058] Step 4: define the pruning ratio, select a specific value in the channel as the weight of each layer channel to judge its contribution to the model, delete the channels in the input layer weight lower than the pruning ratio, and adjust the corresponding channels in the output layer;

[0059] YOLOv4-tiny has up-sampling layer, maximum pooling layer, linear calculation layer and 2 heads. If these layers are changed, the model structure of YOLOv4-tiny will change. Channel pruning is to prune the number of convolution kernels of convolution layer without changing the model structure, thereby reducing the parameters and speeding up the detection speed of the model. Therefore, it is necessary to find the channels with lower contribution to the model for pruning. This paper selects the weight value of the BN (Batch Normalization) layer as the judgment of the contribution of the channel to the model. The calculation formula of the BN layer is as follows:

[0060] y (k) =γ (k) x k +β (k) (1)

[0061] During training, if only ordinary standardization processing is performed, the output distribution of each layer is almost unchanged because the mean and variance are fixed at 0 and 1, respectively. When the γ (k) , β (k) parameters are added, the model can find an optimal output distribution according to continuous learning, and the mean and variance of each layer will change.

[0062] Since β (k) represents the mean of the current layer channel, this paper judges the importance of the channel to the model according to the size of the β (k) parameter. The channels of each layer of the convolution layer are sorted according to the size of the weight value. According to the size of the pruning rate, a weight threshold is found, and the channels with weight values less than the threshold are pruned. The remaining ones are the pruned YOLOv4-tiny model. The structure of the pruned YOLOv4-tiny model is shown in Figure 3 .

[0063] Step 5: retrain the pruned YOLOv4-tiny model to obtain the weight file trained to convergence and load it into the model for testing.

[0064] Because the channel number of each layer of the YOLOv4-tiny model changes, the classification and position information of the model will be deviated, so the pruned model needs to be trained with a small learning rate to update the weight values of each channel, reduce the position information and classification deviation, and after the training is completed, the weight file is saved and loaded into the model for testing to calculate the accuracy, recall rate and model parameter amount of the pruned model;

[0065] Step 6: using the YOLOv4 model as the teacher model of the knowledge distillation algorithm, the vehicle training data set after data augmentation in step 2 is put into the YOLOv4 model for training, and after the training converges, the weight file after the training is completed is saved;

[0066] In the embodiment, because the pruned YOLOv4-tiny model needs to be knowledge distilled, the knowledge information is transferred from the complicated teacher model to a more easily deployed YOLOv4-tiny model, so the YOLOv4 model is selected as the teacher model of the knowledge distillation, the YOLOv4 model structure is complex, but the detection accuracy of the model is high, the YOLOv4-tiny model learns the knowledge information of the YOLOv4 model, thereby improving the detection accuracy of the YOLOv4-tiny model, so the YOLOv4 model needs to be trained, the training data is put into the YOLOv4 model for training, and after the training converges, the weight file is saved to provide the information in the feature map for the knowledge distillation of the following model;

[0067] Step 7: the pruned YOLOv4-tiny model is used as the student model of the knowledge distillation algorithm; the vehicle training data is put into the teacher model and the student model respectively to obtain the feature map; the IoU value of the prior box and the real box is calculated, and the prior box with the IoU value greater than the threshold is calculated by the cross entropy function;

[0068] In the embodiment, as shown in Figure 4 , the YOLOv4 model obtained in step 6 is selected as the teacher model, and the pruned YOLOv4-tiny model in step 5 is selected as the student model; because the YOLOv4 model has three heads, the sizes of the feature maps are 76×76, 38×38 and 19×19, but the student model YOLOv4-tiny needs to be knowledge distilled to have a simple structure with only two heads, and the sizes of the feature maps are 38×38 and 19×19; one head of the YOLOv4 model, that is, the size of the feature map is 76×76, is discarded, and only two heads of the YOLOv4 model are used for knowledge distillation; after the prior box needed for distillation is found, the knowledge distillation of the YOLOv4-tiny model is realized;

[0069] As shown in Figure 5As shown, when using two feature maps to calculate the loss function, unlike traditional object classification-based knowledge distillation, the prior box-based target detection model will construct 3 prior boxes of different sizes on each feature map to determine the target object and the background, resulting in too many negative samples in the prior boxes. Traditional knowledge distillation benefits less from the large model and only focuses on the local area near the real label box to learn knowledge from the local area and calculate the loss function of the two feature maps.

[0070] The training picture will generate a feature map of WxHxK size after passing through the backbone network, where W represents the width of the feature map, H represents the height of the feature map, and K represents the number of feature maps. One-stage detection will generate three prior boxes of different sizes at each position of the feature map. Taking a feature map as an example, 3xWxH prior boxes will be generated. In order to find the prior box near the feature that is close to the real box of the object, this paper uses the method of calculating the IoU between the real box and all prior boxes to estimate the prior box near the real box. The IoU value represents the intersection over union between the real box of the object and the prior box. The larger the IoU value, the larger the overlapping area between the real box and the prior box. Conversely, the IoU value is 0, which means there is no overlapping part. After calculation, a 3xWxH feature matrix will be obtained, and the value in the matrix represents the IoU value between the real box and the prior box at the position. The maximum IoU value M is obtained, and this paper defines a threshold of 0.7M. The prior box with an IoU value greater than this threshold is the local position feature we need. A 3xWxH matrix I is created, and the position feature value is set to 1. Otherwise, it is 0. The IoU calculation is shown below:

[0071]

[0072] where S Truth is the area of the real box, and S box is the area of the prior box.

[0073] The loss function is calculated as follows:

[0074]

[0075] where N p represents the number of 1s in the generated matrix I, which is the number of prior boxes near the position of the real box. W is the width of the feature map, H is the height of the feature map, C is the number of prior boxes of different sizes, s ijc represents the value of the position generated by the student model feature map, and t ijc represents the value of the position generated by the teacher model feature map.

[0076] Step 8: Train the YOLOv4-tiny model using the cross-entropy function in step 7 and the Loss function of YOLOv4-tiny, and save the weight file after training to convergence;

[0077] The overall loss function of the student model is:

[0078] L = L gt + γL i (4)

[0079] Where L gt represents the detection loss during training of the YOLOv4tiny model, γ is the weight factor of the knowledge distillation loss function; the weight factor is selected as 0.01, and L i is the loss function of knowledge distillation obtained in step 7;

[0080] After pruning the YOLOv4-tiny model, perform knowledge distillation, and save the weight file after training is completed;

[0081] Step 9: Test the YOLOv4-tiny model completed by knowledge distillation on the vehicle test data set obtained in step 1 to determine the position coordinates and mark the corresponding positions in the original vehicle test data set;

[0082] Load the weight file into the YOLOv4-tiny model for testing, and use the data set in step 1 for classification to determine the classification and position coordinates of the target object, and mark the corresponding positions in the original test data set; test the accuracy, recall rate, and model parameter amount of the model, as shown in Figure 6 The original picture is shown in Figure 7 , and the detection result of the picture is shown in

Claims

1. A vehicle detection method based on lightweight model compression, characterized in that, The method includes the following steps: Step 1: Collect the vehicle dataset required for model training, and divide the dataset into vehicle training dataset and vehicle test dataset according to the proportions. Step 2: Use data augmentation algorithms to augment and expand the vehicle training dataset, making it twice the size of the original dataset. Step 3: Put the augmented vehicle training dataset from Step 2 into the YOLOv4-tiny model for training. After the training converges, save the weight file after training. Step 4: Set the pruning ratio, select a specific value in the channel as the weight of each channel, judge its contribution to the model, delete the input layer channels whose weights are lower than the pruning ratio, and adjust the output layer channels corresponding to the deleted input layers. Step 5: Retrain the pruned YOLOv4-tiny model to obtain the weight file that has converged, and load it into the model for testing; Step 6: Use the YOLOv4 model as the teacher model for the knowledge distillation algorithm. Put the vehicle training dataset after data augmentation in Step 2 into the YOLOv4 model for training. After the training converges, save the weight file after training. Step 7: Use the pruned YOLOv4-tiny model as the student model for the knowledge distillation algorithm; put the vehicle training data into the teacher model and student model respectively to obtain feature maps; calculate the IoU value between the prior boxes and the ground truth boxes, set a threshold, and calculate the cross-entropy function for the prior boxes with IoU values ​​greater than the threshold. Step 8: Use the cross-entropy function from Step 7 and the loss function of YOLOv4-tiny to sum and train the YOLOv4-tiny model. After training until convergence, save its weight file. Step 9: Use the YOLOv4-tiny model obtained in Step 1 to test the vehicle test dataset, determine the location coordinates, and make corresponding labels in the original vehicle test dataset. The knowledge distillation algorithm used in step 7 to learn feature information from the teacher model is as follows: After the color image passes through the backbone networks of two models, two feature maps of equal size are generated, but the number of channels in the feature maps are not equal. Adaptive feature map convolution is used to make the number of channels in the student model equal to the number of channels in the teacher model through the convolutional layer. The loss function between the two feature maps is then calculated. The specific process is as follows: 1) Introduce the YOLOv4 model trained in step 6 as the teacher model for the knowledge distillation algorithm; 2) Introduce the YOLOv4-tiny model pruned in step 5 as the student model for the knowledge distillation algorithm; 3) After the vehicle training dataset is fed into the two models, two feature maps of the same size are obtained. Adaptive feature map convolution is used, which means that the number of channels in the student model is equal to the number of channels in the teacher model after passing through the convolutional layer. 4) Calculate the loss values ​​of the two feature maps so that the student model can learn the feature information of the teacher model; The loss value is calculated as shown in equation (1): Where, N p The number of 1s in the generated matrix I represents the number of prior boxes whose positions are near the ground truth bounding boxes. W is the width of the feature map, H is the height of the feature map, C is the number of prior boxes of different sizes, and s is the number of prior boxes of different sizes. ijc The value t represents the location of the feature map generated by the student model. ijc The value represents the location of the feature map generated by the teacher's model.

2. The vehicle detection method using lightweight model compression as described in claim 1, characterized in that, The specific process of step 1 is as follows: the vehicle video recorded by the dashcam during the vehicle's operation is used, the sequence in the video is converted into frame-by-frame images as the dataset, the detected objects in the images are labeled, and the images are used as the vehicle training set and the vehicle test set according to the proportion.

3. The vehicle detection method based on lightweight model compression according to claim 1, characterized in that, The specific process of step 2 is as follows: using the Mosaic data augmentation method, four images are mixed by random cropping, random scaling and random arrangement, and the mixed image is used to train the model.

4. The vehicle detection method based on lightweight model compression according to claim 1, characterized in that, The specific value in step 4 is the weight value of the BN layer.

5. The vehicle detection method based on lightweight model compression according to claim 1, characterized in that: The IoU calculation process in step 7 is shown in equation (2): Among them, S Truth S is the area of ​​the true bounding box. box Let be the area of ​​the prior bounding box.

Images