Target detection-oriented model compression method combining iterative pruning and knowledge distillation

Abstract

The application discloses a model compression method for target detection combining iterative pruning and knowledge distillation. The method introduces an assistant network model to relieve the opposite optimization effect caused by the large difference in size between the teacher network and the student network. In the knowledge distillation process, the assistant network is optimized together with the student network, so that the lightweight model after model compression has very good target detection precision and is easy to deploy in mobile terminals or embedded devices to perform target detection tasks. Compared with other integrated pruning model compression methods, the iterative pruning strategy of the method can make the target detection model obtained through final compression have less precision loss on the basis of the same model compression rate.

Description

Technical Field

[0001] This invention is designed for the field of target detection and model compression based on deep neural networks, specifically a model compression method for target detection tasks that combines iterative pruning and knowledge distillation. Background Technology

[0002] With the development of deep learning, the accuracy of object detection models is getting higher and higher, but their model size and number of parameters are also increasing. However, for scenarios that need to perform object detection on mobile devices or embedded devices, excessively large object detection models are not suitable.

[0003] For example, in hazardous industrial scenarios such as underground mines, deep learning technology can assist in manually ensuring workers wear their helmets, preventing them from entering without them. This provides a safeguard for worker safety. Currently, various object detection networks are emerging in academia, achieving excellent detection results in object detection scenarios, such as YOLOv3 and YOLOv5. However, since most object detection networks are based on convolutional neural networks, their network structures are complex and have a large number of parameters. In actual industrial production environments, mobile or embedded devices have limited memory and computing power, posing significant challenges to the deployment of large-scale convolutional object detection networks. Summary of the Invention

[0004] To address the shortcomings of the existing technologies, the present invention aims to provide a model compression method for object detection that combines iterative pruning and knowledge distillation. This method comprises two main steps: iterative pruning followed by knowledge distillation to improve detection accuracy. The object detection network processed by this method exhibits lightweight structure, high detection accuracy, and ease of deployment.

[0005] The specific technical solution adopted in this invention is as follows:

[0006] A model compression method for object detection based on joint iterative pruning and knowledge distillation includes the following steps:

[0007] S1. Obtain the image dataset used for model training in the object detection task, and the image samples in the image dataset are pre-labeled with the category and location of the target to be detected;

[0008] S2. Use the image dataset to perform iterative training on the teacher network for a specified number of rounds;

[0009] S3, performing specified rounds of iterative training on the teaching network using the image dataset, and performing sparse training on the teaching network; the teaching network and the student network are both object detection networks, and the student network has a simpler network structure and fewer model parameters than the teaching network;

[0010] S4, performing iterative pruning on the student network after the sparse training, thereby compressing the model; and during the iterative pruning, retraining the model after each pruning to reduce the loss of accuracy, and obtaining the student network after completing all the iterative rounds;

[0011] S5, using the trained teaching network in S2, the trained student network in S3, and the obtained student network in S4 to perform knowledge distillation; and during the knowledge distillation, using the knowledge provided by the teaching network to continuously optimize the student network and the student network, and finally obtaining the student network model for deployment in industrial mobile devices or embedded devices to perform object detection tasks.

[0012] Preferably, the teaching network adopts a YOLOV5m network.

[0013] Preferably, the student network adopts a YOLOV5s network.

[0014] Preferably, the pruning operation in S4 uses a channel pruning strategy, and an iterative pruning and retraining process is used to implement the pruning operation; before the iteration, the total model pruning ratio P is divided according to the total number of iterations M, so that the pruning ratio p of each iteration satisfies (1-p) M deviation from 1-P is less than a threshold value, and in each iteration, the student network is pruned according to the pruning ratio p, and after the pruning is completed, the student network is retrained once; after M iterations, the student network is obtained.

[0015] Preferably, in S5, the student network and the student network obtain the knowledge of the teaching network model during the knowledge distillation process, and are optimized, and the loss functions for optimizing the student network and the student network are as follows:

[0016] The loss function of the student network is:

[0017]

[0018] wherein KL represents a KL distance calculation function, CE represents a cross-entropy loss calculation function, N represents the total number of samples output by the network, N p represents the number of positive samples output by the network, y as i represents the classification result output by the student network for the i-th sample, y gt iy represents the target classification label of the i th sample t j represents the classification result output by the teacher network for the i th sample, and represents the distillation loss weight

[0019] The loss function of the student network is the sum of the classification loss and the regression loss , which is in the form of:

[0020]

[0021]

[0022]

[0023]

[0024] wherein, s i represents the classification result output by the student network for the i th sample, and represents the regression result output by the student network for the i th sample s i represents the prediction box coordinate result output by the student network for the i th sample, and represents the prediction box coordinate result output by the student network for the i th sample gt i represents the labeled box coordinate of the i th sample, and represents the labeled box coordinate of the i th sample as j represents the prediction box result output by the teaching network for the i th sample, and represents the prediction box result output by the teaching network for the i th sample t j represents the prediction box result output by the teaching network for the i th sample, and represents the prediction box result output by the teaching network for the i th sample, wherein represents the distillation loss weight, and represents a threshold value; L(r s j represents the CIOU loss of r as j . s j as j .

[0025] As a preferred, the threshold value is set to 0.7.

[0026] As a preferred, the distillation loss weight is set to 10.

[0027] Compared with the prior art, the present application has the beneficial effects of:

[0028] 1. Compared with other integrated pruning model compression methods, the iterative pruning strategy of the present application can make the final compressed target detection model have less precision loss on the basis of the same model compression rate.

[0029] ​2、The method introduces an assistant network model, alleviates the opposite optimization effect caused by the large difference in size between the teacher network and the student network, and optimizes the assistant network and the student network during knowledge distillation, so that the lightweight model after model compression has very good target detection detection accuracy and is easy to deploy in mobile terminals or embedded devices to perform target detection tasks. BRIEF DESCRIPTION OF DRAWINGS

[0030] Figure 1 Flowchart of the method of the present application

[0031] Figure 2 Effect diagram of iterative pruning of the present application

[0032] Figure 3 Overall framework diagram of the present application combining iterative pruning and knowledge distillation based on the assistant network

[0033] Figure 4 Prediction effect diagram in an embodiment of the present application. DETAILED DESCRIPTION

[0034] In order to make the above-mentioned purposes, features and advantages of the present application more obvious and easy to understand, the specific embodiments of the present application will be described in detail below with reference to the drawings. In the following description, many specific details are set forth in order to provide a thorough understanding of the present application. However, the present application can be implemented in many other ways different from those described herein, and those skilled in the art can make similar improvements without departing from the scope of the present application, so the present application is not limited to the specific embodiments disclosed below. The technical features in each embodiment of the present application can be combined accordingly without conflict.

[0035] In a preferred embodiment of the present application, a target detection-oriented model compression method combining iterative pruning and knowledge distillation is provided, as shown in Figure 1 which will be described in detail as follows:

[0036] Step 1, obtaining an image dataset used for model training in a target detection task, and each image sample in the image dataset is pre-labeled with the class and position of the target to be detected.

[0037] In this embodiment, the image data corresponding to the target to be detected can be obtained by using a camera to record, using a web crawler to crawl, etc. according to the task requirements, and the dataset is cleaned and labeled in advance, mainly labeling the class and corresponding position of the target to be detected in the dataset. Of course, the image dataset can be collected by oneself, or an existing labeled dataset of a target detection task can be obtained. The image dataset can be divided into a training set and a validation set to meet the model training and testing requirements.

[0038] Step 2, according to the task requirements, a teacher network model is designed, and the teacher network model is iteratively trained for a specified number of rounds using the above image dataset.

[0039] Since the teacher network model needs to provide corresponding detection information in the knowledge distillation process, it has a high requirement for its accuracy, but there is no limit to its model complexity and parameter size, so some network models with complex structure, large parameter size and high detection accuracy can be used as the teacher network model. In this embodiment, YOLOV5m can be used as the teacher network model, and the training data is input for 100 epochs of training, which is used for the subsequent knowledge distillation stage.

[0040] Step 3, design a teaching assistant network model, and iteratively train the teaching assistant network for a specified number of rounds using the image dataset, and then perform sparse training on the teaching assistant network.

[0041] The role of the teaching assistant network model introduced in the present application is mainly to alleviate the problem of poor effect in the knowledge distillation process caused by the large gap between the teacher network model and the student network model. In the knowledge distillation process, the teaching assistant network is used to transition the knowledge of the complex structure teacher network model to the relatively simple student network model, so as to optimize the detection accuracy of the student network model. At the same time, the teaching assistant network model also needs to be sparsely trained. The purpose of sparse training is to make the model parameters sparse, which facilitates the subsequent iterative channel pruning operation.

[0042] The teacher network and the teaching assistant network are both object detection networks, and the teaching assistant network should have a simpler network structure and fewer model parameters than the teacher network. For the teaching assistant network model, YOLOV5s with relatively simple structure and relatively small parameter size can be used as the teaching assistant network model in this embodiment, and the training data is input for 100 epochs of training. After training, the model is sparsely trained for 100 epochs, which is used for subsequent pruning processing.

[0043] Step 4, the pruned teaching assistant network is iteratively pruned to compress the model, and the pruned model is retrained after each pruning to reduce its accuracy loss, and the student network is obtained after all iterations are completed.

[0044] In this embodiment, the pruning operation can use the channel pruning strategy. At the same time, in order to further optimize the channel pruning strategy of the YOLOV5 network model, the original integrated pruning and retraining process is optimized into an iterative pruning and retraining process, and a minimum channel retention ratio is set for each iteration. Specifically, before iteration, the total model pruning ratio P is decomposed according to the total number of iterations M, so that the pruning ratio p of each iteration satisfies (1-p)M If the deviation from 1-P is less than a threshold, the teaching network is pruned according to a pruning ratio p in each iteration. After pruning, the teaching network is retrained. After M iterations, the student network is obtained. For example, if the model needs to be pruned by a ratio of P = 75%, the remaining model size should theoretically be 25% of the original model. After pruning, retraining is used to reduce accuracy loss. This process is broken down so that the pruning ratio p = 30% in each iteration. In each iteration, only 30% of the model is pruned, and after pruning, the model is retrained. This process is repeated four times. Figure 2 As shown, with each iteration, the model size is further reduced, so the final detection model size is theoretically (1-30%) of the original model. 4 ≈24%. This yields a model compression effect that is almost identical to the original integrated pruning. However, since the pruning ratio is smaller in each iteration and there is a retraining process in each iteration, the iterative pruning and retraining process results in a better final pruning effect compared to integrated pruning.

[0045] Step 5: Use the teacher network trained in S2, the teaching assistant network trained in S3, and the student network obtained in S4 to perform knowledge distillation. During the knowledge distillation process, the knowledge provided by the teacher network is used to continuously optimize the teaching assistant network and the student network. The final student network model is significantly compressed and its accuracy is also greatly improved. Therefore, the final knowledge distillation student network model is used to deploy in industrial mobile devices or embedded devices to perform target detection tasks.

[0046] In this invention, during the knowledge distillation process, both the teaching assistant network model and the student network model simultaneously acquire knowledge from the teacher network model, thereby achieving optimization. For example... Figure 3 The diagram illustrates the learning process of the teaching assistant network model and the student network model under the knowledge assistance provided by the teacher network model during knowledge distillation. The core of this process is calculating the loss functions of the teaching assistant network and the student network. The optimization of the loss functions of the teaching assistant network and the student network in this embodiment is described in detail below:

[0047] Step 5-1, for the assistant network, since its network has been trained and has good detection effect, only the classification loss is considered in the distillation training process. The classification loss is divided into two categories. One is caused by the difference with the teacher network model, which is measured by KL distance (Kullback-Leibler Divergence), and the other is caused by the difference with the label, which is measured by cross entropy loss CE (Cross Entropy). The loss function can be expressed as:

[0048]

[0049] Where N represents the total number of network output samples, N p represents the number of positive samples output by the network, y as i represents the classification result output by the assistant network for the ith sample, y gt i represents the target classification label of the ith sample, y as j represents the classification result of the assistant network for the ith sample, y t j represents the classification result output by the teacher network for the ith sample, β represents the distillation loss weight. In the distillation training process, the distillation loss weight needs to be larger, otherwise there is no distillation effect. Therefore, in this embodiment, the value of β can be set to 10.

[0050] Step 5-2, for the student network, the loss function is composed of two parts: classification loss and regression loss. For the classification loss, the difference between the output of the student network and the label needs to be considered, which is measured by cross entropy loss CE (Cross Entropy). At the same time, the knowledge distillation loss between the assistant network and the student network needs to be introduced, which is measured by KL distance (Kullback-Leibler Divergence). For the regression loss, the difference between the predicted frame output by the student network and the labeled frame needs to be considered, which is measured by CIoU_loss. At the same time, when the IoU between the coordinate information of the assistant network and the teacher network is greater than a specified threshold δ (generally set to 0.7), the coordinate difference between the assistant network and the student network also needs to be introduced, which is also measured by CIoU_loss. In summary, in the distillation training process, the loss function of the student network is:

[0051]

[0052]

[0053]

[0054]

[0055] Wherein, N represents the total number of samples of network output, N p represents the number of positive samples of network output, y s i represents the classification result of the student network output for the i-th sample, y gt i represents the target classification label of the i-th sample, y as j represents the classification result of the i-th sample of the teaching assistant network, r s i represents the predicted box coordinate result of the student network output for the i-th sample, r gt i represents the annotation box coordinate of the i-th sample, r as j represents the predicted box result of the i-th sample output by the teaching assistant network, r t j represents the predicted box result of the i-th sample output by the teacher network, β represents the distillation loss weight, and in the distillation training process, the distillation loss weight needs to be larger, otherwise there is no distillation effect, so the value of β in the embodiment can be set to 10. L(r s j , r as j represents the coordinate difference loss between the teaching assistant network and the student network, and CIoU_loss is used to measure, that is, the CIOU loss of r s j , r as j .

[0056] Figure 4 In an example of the present application, the above-mentioned student network is applied in the task of detecting workers wearing safety helmets, and the prediction effect diagram of the input actual image obtained is as follows, wherein the WithHelmet category represents that the worker wears a safety helmet, and the WithoutHelmet category represents that the worker does not wear a safety helmet. As can be seen from the figure, in the extremely complex background scene, the target is still accurately detected. Therefore, the lightweight model after compression of the above-mentioned model has fewer model parameters, but still maintains good detection accuracy, and can be deployed in mobile terminals or embedded devices. For example, the above-mentioned student network model can be deployed in a worker safety helmet wearing detection system under the mine to improve the safety guarantee of the worker working under the mine.

[0057] The above-described embodiments are only the preferred ones of the present application, and are not intended to limit the present application. Those skilled in the art can make various changes and modifications without departing from the spirit and scope of the present application. Therefore, any technical solutions obtained by equivalent replacement or equivalent transformation shall fall within the protection scope of the present application.

Claims

1. A target detection-oriented model compression method combining iterative pruning and knowledge distillation, characterized in that, The method comprises the following steps: S1, obtaining an image dataset for model training in a target detection task, and each image sample in the image dataset is pre-labeled with a category and a position of a target to be detected; S2, performing a specified number of iterations of training on a teacher network using the image dataset; S3, performing a specified number of iterations of training on a teaching assistant network using the image dataset, and then performing sparse training on the teaching assistant network; the teacher network and the teaching assistant network are both target detection networks, and the teaching assistant network has a simpler network structure and fewer model parameters than the teacher network; S4, performing pruning operations on the teaching assistant network after sparse training to compress the model; during the iterative pruning process, the model after each pruning is retrained to reduce the loss of accuracy, and the student network is obtained after all iterations are completed; S5, using the teacher network trained in S2, the teaching assistant network trained in S3, and the student network obtained in S4 to perform knowledge distillation; during the knowledge distillation process, the knowledge provided by the teacher network is used to continuously optimize the teaching assistant network and the student network, and finally the student network model obtained through knowledge distillation is used to perform a target detection task on an industrial mobile device or an embedded device; In S5, the teaching assistant network and the student network obtain the knowledge of the teacher network model during the knowledge distillation process to obtain optimization, and the loss functions for optimizing the teaching assistant network and the student network are as follows: The loss function of the teaching assistant network is as follows: wherein, denotes a KL distance computation function, denotes a cross-entropy loss computation function, denotes the total number of samples output by the network, denotes the number of positive samples output by the network, denotes the classification result output by the assistant network for the i-th sample, denotes the target classification label of the i-th sample, denotes the classification result output by the teacher network for the i-th sample, denotes a distillation loss weight; The loss function for the student network is the sum of a classification loss and a regression loss in the form of: in, This represents the classification result output by the student network for the i-th sample. This represents the predicted bounding box coordinates output by the student network for the i-th sample. This represents the coordinates of the bounding box for the i-th sample. This represents the predicted bounding box output by the teaching assistant network for the i-th sample. This represents the predicted bounding box output by the teacher network for the i-th sample. Indicates the weight of distillation loss. Indicates the threshold; Indicates calculation CIOU Loss.

2. The target detection oriented model compression method of joint iterative pruning and knowledge distillation according to claim 1, characterized in that, The teacher network adopts a YOLOV5m network.

3. The target detection oriented model compression method of joint iterative pruning and knowledge distillation according to claim 1, characterized in that, The teaching assistant network adopts a YOLOV5s network.

4. The target detection oriented model compression method of joint iterative pruning and knowledge distillation according to claim 1, characterized in that, The pruning operation in the S4 uses a channel pruning strategy and is implemented by an iterative pruning and retraining process; before iteration, the total model pruning ratio P is decomposed according to the total iteration number M, so that the pruning ratio p of each iteration satisfies (1-p) M The deviation from 1-P is less than a threshold value, and in each iteration, the teaching network is pruned according to the pruning ratio p, and after the pruning is completed, the teaching network is retrained once; after M iteration rounds are completed, the student network is obtained.

5. The target detection oriented model compression method of joint iterative pruning and knowledge distillation according to claim 1, characterized in that, the threshold value is set to 0.

7.

6. The target detection oriented model compression method of joint iterative pruning and knowledge distillation according to claim 1, characterized in that, the distillation loss weight was set to 10.

7. A target detection method characterized by, The student network model obtained by the model compression method according to any one of claims 1-6 is used to perform a target detection task on a collected image.

Images