Target detection model compression method based on step-by-step knowledge distillation

Abstract

The invention discloses a step-by-step knowledge distillation-based target detection model compression method. The method comprises the following steps of making a target detection data set; selecting Drone-YOLO as a teacher model, and training the teacher model by using the target detection data set; yOLOv8 is selected as a student model, the student model is provided with five model structures including YOLOv8x, YOLOv81, YOLOv8m, YOLOv8s and YOLOv8n, the five model structures of the student model are respectively modified to have the same output as the teacher model, then a loss function L is defined, the five modified student model structures are sequentially distilled step by step through the loss function L, and step-by-step transmission of knowledge is achieved; according to the application scene, an optimal model structure in the student models is selected for deployment, and it is ensured that an efficient target detection function is achieved in the target application scene. Through a step-by-step distillation strategy, the current model can inherit and optimize the performance from the previous model, and the performance loss caused by model compression is reduced.

Description

Technical Field

[0001] The present invention belongs to the technical field of deep learning and computer vision, and relates to a target detection model compression method based on step-by-step knowledge distillation. Background Art

[0002] In the field of target detection, the YOLO (You Only Look Once) series of models are popular for their real-time processing capabilities and high efficiency; the Yolov8 series has five models: x, l, m, s, and n, with their computational complexity and accuracy ranked from high to low. Currently, the Yolov8 model is good enough in processing image data taken from a normal perspective. However, most drones are shot from a bird's-eye view, and the image quality is easily affected by the lighting, atmospheric environment, and motion posture of the drone during shooting. The size of the target is also mostly small. Therefore, the effect of directly using the Yolov8 series model to train the business data set shot by the drone is limited. It is often necessary to first train and tune a large model to an ideal level, and then use model compression technology to migrate its feature extraction and fusion capabilities to the model that is ultimately used on the hardware.

[0003] Current model compression technologies mainly include network pruning, quantization, low-rank decomposition, and knowledge distillation. When performing model compression, especially using knowledge distillation, performance loss often occurs because the student model may not be able to fully capture all the information in the teacher model, especially in highly compressed situations; the traditional single-step distillation process usually only involves the migration from a large teacher model to a small student model. A one-time large-scale compression may lead to the loss of key information and fail to efficiently utilize the knowledge of the teacher model; moreover, existing distillation strategies are often fixed and lack flexibility, making it difficult to adjust according to the specific needs of different application scenarios, such as requiring different degrees of model compression under different hardware conditions; the direct conversion from a large model to a small model does not make full use of the intermediate model for gradual performance adjustment and optimization. Summary of the invention

[0004] In view of this, the purpose of the present invention is to provide a target detection model compression method based on step-by-step knowledge distillation to solve the technical problems mentioned in the prior art.

[0005] The object detection model compression method based on step-by-step knowledge distillation includes the following steps:

[0006] S1, create target detection dataset;

[0007] S2. Select Drone-YOLO as the teacher model and train the teacher model using the target detection dataset;

[0008] YOLOv8 is selected as a student model, and the student model has five model structures, namely, YOLOv8x, YOLOv8l, YOLOv8m, YOLOv8s and YOLOv8n. The five model structures of the student model are respectively modified to have the same output as the teacher model, and the modified five student model structures are respectively trained using the target detection dataset;

[0009] S3. Define the loss function L;

[0010] L = (1-α) × L CE (s,y)+α×T 2 ×L KL

[0011] Among them, L CE is the cross entropy loss of the student model for the hard label, which is used to measure the error between the student model output and the true label, s is the student model output, and y is the true label; L KL is the KL divergence, which measures the similarity between the soft label and the output of the student model; T is the temperature parameter used to soften the output probability distribution; α is a balance parameter used to control the weight of the two losses;

[0012] S4, distilling the modified five student model structures step by step in sequence through the loss function L to achieve step-by-step transfer of knowledge;

[0013] S5. According to the application scenario, select the best model structure in the student model for deployment to ensure efficient target detection function in the target application scenario.

[0014] Optionally, in S1, an image is captured and annotated with an image annotation tool to obtain hard labels, and the annotated pedestrians and vehicles are used as target detection subjects to generate a target detection data set.

[0015] Optionally, in S3, the Head network of the teacher model has four detection heads, and the four detection heads of the teacher model respectively output feature maps of sizes of B*(nc+16*4)*(160*160), B*(nc+16*4)*(80*80), B*(nc+16*4)*(40*40) and B*(nc+16*4)*(20*20) from top to bottom;

[0016] The Head network of the student model has three detection heads, and the three detection heads of the student model output feature maps of sizes B*(nc+16*4)*(80*80), B*(nc+16*4)*(40*40) and B*(nc+16*4)*(20*20) from top to bottom respectively;

[0017] Among them, B is the batch_size parameter, nc is the number of recognition categories;

[0018] Add a set of structural layers in the Head network of the student model, wherein;

[0019] By adjusting the parameter sizes of d, w and r of the student model, the number of some layers in each model structure is adjusted to change the number of input feature map channels of some layers of the student model so that the output of the student model is the same as that of the teacher model.

[0020] Optionally, in S3, the soft label is obtained by predicting the target detection dataset through the teacher model.

[0021] Optionally, in S4, the method of performing step-by-step distillation on the student model structure is:

[0022] ① Use Drone-YOLO as the teacher model and distill the modified student model YOLOv8x through the loss function L;

[0023] ② Use the distilled student model YOLOv8x as the teacher model, and distill the modified student model YOLOv8l through the loss function L;

[0024] ③ Use the distilled student model YOLOv8l as the teacher model, and distill the modified student model YOLOv8m through the loss function L;

[0025] ④ Use the distilled student model YOLOv8m as the teacher model, and distill the modified student model YOLOv8s through the loss function L;

[0026] ⑤ Use the distilled student model YOLOv8s as the teacher model, and distill the modified student model YOLOv8n through the loss function L.

[0027] Optionally, the distillation method is:

[0028] The student model to be distilled is trained by the loss function L;

[0029] Performing performance evaluation on the trained student model to obtain a performance evaluation result;

[0030] The distillation parameters were optimized based on the performance evaluation results.

[0031] Optionally, the objects of the performance evaluation are the accuracy, detection speed and resource consumption of the student model.

[0032] Optionally, the distillation parameters include at least a temperature parameter and a weight coefficient.

[0033] The beneficial effects that the present invention can produce include:

[0034] 1) The target detection model compression method based on step-by-step knowledge distillation provided by the present invention enables the current model to inherit and optimize the performance from the previous model through a step-by-step distillation strategy, thereby reducing the performance loss caused by model compression.

[0035] 2) The target detection model compression method based on step-by-step knowledge distillation provided by the present invention gradually improves the capabilities of the student model by adopting a multi-stage distillation method. Each stage can further utilize the knowledge learned in the previous stage, thereby improving the overall distillation efficiency.

[0036] 3) The target detection model compression method based on step-by-step knowledge distillation provided by the present invention allows the selection of models at different stages according to different deployment requirements, providing more flexible model selection for different application scenarios to adapt to different performance and resource requirements.

[0037] 4) The target detection model compression method based on step-by-step knowledge distillation provided by the present invention can optimize the model at each stage on the basis of inheriting the performance of the previous stage, gradually approaching the balance point between the optimal performance and the compression ratio. BRIEF DESCRIPTION OF THE DRAWINGS

[0038] Figure 1 It is a flow chart of the target detection model compression method based on step-by-step knowledge distillation of the present invention;

[0039] Figure 2 Schematic diagram of the Drone-YOLO network structure of the present invention;

[0040] Figure 3 This is a schematic diagram of the modified YOLOv8 network structure of the present invention. DETAILED DESCRIPTION

[0041] The present invention is described in detail below with reference to the accompanying drawings and embodiments, but the present invention is not limited to these embodiments.

[0042] The object detection model compression method based on step-by-step knowledge distillation includes the following steps:

[0043] Step 1: Create a target detection dataset. In this step, capture images, such as using a drone to take aerial photos, and use image annotation tools to annotate pedestrians and vehicles in the images to obtain hard labels. The annotated pedestrians and vehicles are used as target detection subjects to generate a target detection dataset.

[0044] Step 2. Select Drone-YOLO as the teacher model and use the target detection dataset to train the teacher model to obtain the initial weight file Drone-YOLO.pt; select YOLOv8 as the student model. The student model has five model structures: YOLOv8x, YOLOv8l, YOLOv8m, YOLOv8s, and YOLOv8n. Modify the five model structures of the student model to have the same output as the teacher model, and use the target detection dataset to train the modified five student model structures respectively to obtain YOLOv8x.pt, YOLOv8l.pt, YOLOv8m.pt, YOLOv8s.pt, and YOLOv8n.pt files.

[0045] In the above, the Head network of the teacher model has four detection heads, and the four detection heads of the teacher model output feature maps with sizes from top to bottom of B*(nc+16*4)*(160*160), B*(nc+16*4)*(80*80), B*(nc+16*4)*(40*40) and B*(nc+16*4)*(20*20); the Head network of the student model has three detection heads, and the three detection heads of the student model output feature maps with sizes from top to bottom of B*(nc+16*4)*(80*80), B*(nc+16*4)*(40*40) and B*(nc+16*4)*(20*20). 40) and B*(nc+16*4)*(20*20) feature maps; wherein B is the batch_size parameter, nc is the number of recognition categories, which depends on the target detection data set. In the present embodiment, nc is set to 2; a group of structural layers are added to the Head network of the student model, wherein the number of some layers in each model structure is adjusted (such as the number of repetitions of the BootleNeck block in the C2f layer) by adjusting the parameter sizes of d, w and r of the student model to change the number of input feature map channels of some layers of the student model so that the output of the student model is the same as that of the teacher model.

[0046] Specifically, Figure 3 As shown in the figure, when the parameters of d, w and r of the student model are the default values ​​1, the model structure of the student model is Yolov8l, and the adjustment process of the Yolov8l model structure is:

[0047] ① Add a set of Upsample layers (numbered 16), Concat layers (numbered 17) and C2f layers (numbered 18) to the original structure layer of the student model, and then pass the output of the C2f layer numbered 15 to the Upsample layer numbered 16 for upsampling to obtain a feature map of size B*160*160*256;

[0048] ② The feature map output by the Upsample layer numbered 16 is fused with the feature map output by the C2f layer numbered 2 at the channel layer, and the fused feature map is input into the C2f layer numbered 18 to obtain a feature map of B*160*160*128, and the feature map of B*160*160*128 is used as the input of the newly added detection head; the other branch performs 3*3 convolution on the fused feature map to generate a feature map of B*80*80*128, and fuses it with the feature map output by the C2f layer numbered 15 at the channel layer, and inputs the fused feature map into the C2f layer numbered 21 to obtain a feature map of B*80*80*256, and the feature map of B*80*80*256 is used as the input of the second detection head and the input of the convolution layer numbered 22;

[0049] ③ The feature map output by the C2f layer numbered 12 is fused with the feature map output by the convolution layer numbered 22 at the channel layer, and the fused feature map is input into the C2f layer numbered 24 to obtain a feature map of B*40*40*512, and the feature map of B*40*40*512 is used as the input of the third detection head;

[0050] ④Fuse the feature map output by the convolution layer numbered 25 with the feature map output by the SPPF layer numbered 9 at the channel layer, input the fused feature map into the C2f layer numbered 27 to obtain a feature map of B*20*20*512, and use the feature map of B*20*20*512 as the input of the fourth detection head.

[0051] Step 3: Define the loss function L;

[0052] L = (1-α) × L CE (s,y)+α×T 2 ×L KL

[0053] Among them, L CE is the cross entropy loss of the student model for the hard label, which is used to measure the error between the student model output and the true label, s is the student model output, and y is the true label; L KL is the KL divergence, which measures the similarity between the soft label and the output of the student model; T is the temperature parameter used to soften the output probability distribution; α is a balancing parameter used to control the weights of the two losses; specifically, the soft label is predicted by the teacher model on the target detection dataset.

[0054] Step 4: Use the loss function L to distill the modified five student model structures step by step to achieve step-by-step transfer of knowledge.

[0055] Specifically, the method of performing step-by-step distillation on the student model structure is:

[0056] ① Use Drone-YOLO as the teacher model and distill the modified student model YOLOv8x through the loss function L;

[0057] ②Use the distilled student model YOLOv8x as the teacher model, and distill the modified student model YOLOv8l through the loss function L;

[0058] ③Use the distilled student model YOLOv8l as the teacher model, and distill the modified student model YOLOv8m through the loss function L;

[0059] ④Use the distilled student model YOLOv8m as the teacher model, and distill the modified student model YOLOv8s through the loss function L;

[0060] ⑤ Use the distilled student model YOLOv8s as the teacher model, and distill the modified student model YOLOv8n through the loss function L.

[0061] In the above, the distillation method is: the student model to be distilled is trained through the loss function L, and then the performance of the trained student model is evaluated, such as performing data statistical analysis on the accuracy, detection speed and resource consumption of each student model structure during the training process, and then sorting the performance parameters of each student model structure according to the discreteness according to the analysis results, so as to obtain the performance evaluation results, and then optimizing the distillation parameters based on the performance evaluation results, such as using the data iterative update method, so that the student model structure of each stage can optimize the distillation parameters on the basis of inheriting the performance of the previous stage, and gradually approach the balance point of the optimal performance and compression ratio, wherein the distillation parameters include at least temperature parameters and weight coefficients.

[0062] Step 5. According to the application scenario, select the best model structure in the student model for deployment to ensure efficient target detection in the target application scenario; for example, use the improved YOLOv8n model structure as the target detection model for deployment on resource-constrained drones or mobile devices. Through the above step-by-step distillation strategy, the Drone-YOLO teacher model progressively transfers knowledge through the five levels of YOLOv8 series intermediate student models, which can gradually refine and strengthen the learning ability of the student model, so that each level of the student model structure can fully utilize the knowledge of the previous level of the student model structure, which not only improves the compression efficiency and deployment flexibility of the target detection model, but also maintains the optimization of the target detection model performance, so that the target detection model can still maintain efficient operation in a resource-constrained environment.

[0063] The above are only several embodiments of the present invention and are not intended to limit the present invention in any form. Although the present invention is disclosed as above in the form of a preferred embodiment, it is not intended to limit the present invention. Any technician familiar with the profession, without departing from the scope of the technical solution of the present invention, using the above disclosed technical content to make slight changes or modifications are equivalent to equivalent implementation cases and fall within the scope of the technical solution.

Claims

1. The object detection model compression method based on step-by-step knowledge distillation is characterized by: The following steps are involved: S1, create target detection dataset; S2. Select Drone-YOLO as the teacher model, and use the target detection dataset to train the teacher model; YOLOv8 is selected as a student model, and the student model has five model structures, namely, YOLOv8x, YOLOv8l, YOLOv8m, YOLOv8s and YOLOv8n. The five model structures of the student model are respectively modified to have the same output as the teacher model, and the modified five student model structures are respectively trained using the target detection dataset; S3. Define the loss function L; L=(1-α)×L CE (s,y)+α×T 2 ×L KL Among them, L CE is the cross entropy loss of the student model for the hard label, which is used to measure the error between the student model output and the true label, s is the student model output, and y is the true label; L KL is the KL divergence, which measures the similarity between the soft labels and the output of the student model; T is a temperature parameter used to soften the output probability distribution; α is a balancing parameter used to control the weights of the two losses; S4, distilling the modified five student model structures step by step in sequence through the loss function L to achieve step-by-step transfer of knowledge; S5. According to the application scenario, select the best model structure in the student model for deployment to ensure efficient target detection function in the target application scenario.

2. The target detection model compression method based on step-by-step knowledge distillation according to claim 1 is characterized in that: In S1, an image is captured and annotated with an image annotation tool to obtain hard labels. The annotated pedestrians and vehicles are used as target detection subjects to generate a target detection data set.

3. The object detection model compression method based on step-by-step knowledge distillation according to claim 1 is characterized in that: In S3, the Head network of the teacher model has four detection heads, and the four detection heads of the teacher model output feature maps of sizes B*(nc+16*4)*(160*160), B*(nc+16*4)*(80*80), B*(nc+16*4)*(40*40) and B*(nc+16*4)*(20*20) from top to bottom respectively; The Head network of the student model has three detection heads, and the three detection heads of the student model output feature maps of sizes B*(nc+16*4)*(80*80), B*(nc+16*4)*(40*40) and B*(nc+16*4)*(20*20) from top to bottom respectively; Among them, B is the batch_size parameter, nc is the number of recognition categories; Add a set of structural layers in the Head network of the student model, wherein; By adjusting the parameter sizes of d, w and r of the student model, the number of some layers in each model structure is adjusted to change the number of input feature map channels of some layers of the student model so that the output of the student model is the same as that of the teacher model.

4. The object detection model compression method based on step-by-step knowledge distillation according to claim 1 is characterized in that: In S3, the soft label is obtained by predicting the target detection dataset through the teacher model.

5. The object detection model compression method based on step-by-step knowledge distillation according to claim 1 is characterized in that: In S4, the method of performing step-by-step distillation on the student model structure is: ① Use Drone-YOLO as the teacher model and distill the modified student model YOLOv8x through the loss function L; ② Use the distilled student model YOLOv8x as the teacher model, and distill the modified student model YOLOv8l through the loss function L; ③ Use the distilled student model YOLOv8l as the teacher model, and distill the modified student model YOLOv8m through the loss function L; ④ Use the distilled student model YOLOv8m as the teacher model, and distill the modified student model YOLOv8s through the loss function L; ⑤ Use the distilled student model YOLOv8s as the teacher model, and distill the modified student model YOLOv8n through the loss function L.

6. The object detection model compression method based on step-by-step knowledge distillation according to claim 5 is characterized in that: The distillation method is: The student model to be distilled is trained by the loss function L; Performing performance evaluation on the trained student model to obtain a performance evaluation result; The distillation parameters were optimized based on the performance evaluation results.

7. The object detection model compression method based on step-by-step knowledge distillation according to claim 6, characterized in that: The objects of the performance evaluation are the accuracy, detection speed and resource consumption of the student model.

8. The object detection model compression method based on step-by-step knowledge distillation according to claim 6, characterized in that: The distillation parameters at least include a temperature parameter and a weight coefficient.

Images