An object detection network channel pruning method based on information gain

Abstract

This invention discloses a channel pruning method for object detection networks based on information gain, comprising the following steps: constructing a training dataset and a test dataset; training the object detection network on the training dataset until convergence to obtain an initial unpruned teacher model; calculating the channel information gain values ​​of all BN layers to be pruned in the teacher model; presetting a pruning ratio, and pruning the network channel layers in the object detection network according to the preset pruning ratio to obtain a student model; modifying the training loss function, and using the teacher model to perform output knowledge distillation on the student model to achieve fine-tuning. This method does not require sparse training, the pruning process is simple and easy to implement, and achieves the effect of effectively compressing the object detection network under the premise of low loss accuracy.

Description

Technical Field

[0001] This invention relates to the field of target detection technology, and in particular to a target detection network channel pruning method based on information gain. Background Technology

[0002] Object detection is the task of classifying and locating targets in images or videos. It has broad application prospects and value, and can be directly applied to many practical scenarios such as intelligent transportation, industrial equipment inspection, and defect detection. Taking the application of object detection in the transportation field as an example, in traffic monitoring and security, object detection algorithms can be used to analyze surveillance video streams and detect traffic violations on the road in real time, such as driving against traffic, running red lights, and crossing lane lines, thereby improving monitoring efficiency and traffic safety. In intelligent traffic control, object detection technology can be used to automatically sense traffic flow, realizing intelligent traffic control and scheduling, such as traffic light control, priority passage, and congestion prediction.

[0003] Because deep neural networks have a large number of parameters, some of which contribute little to the network model's output are considered redundant. These parameters can be pruned to reduce the computational cost and parameter count of the network model. Most current channel pruning schemes for object detection networks require initial sparsity training, necessitating multiple parameter fine-tuning steps to achieve good sparsity results. Channel pruning methods often rely on traditional retraining to fine-tune the pruned network and restore accuracy. However, since retraining has limited accuracy recovery capabilities, knowledge distillation is employed to fine-tune the pruned network. Knowledge distillation is a method that transfers knowledge from a pre-trained, complex teacher model to a student model; therefore, during deployment, a smaller student model can replace a larger teacher model.

[0004] Most existing fine-tuning schemes for object detection networks use output knowledge distillation, which only design classification distillation loss and coordinate regression distillation loss, ignoring the fact that a large amount of background information output by the teacher model can affect the bounding box coordinate regression training problem of the student model. Summary of the Invention

[0005] This invention overcomes the shortcomings of the prior art and provides a target detection network channel pruning method based on information gain. Its important purpose is to reduce the number of parameters and computational load of the target detection network while minimizing the loss of detection accuracy.

[0006] To achieve the above objectives, the first aspect of this invention provides a channel pruning method for target detection networks based on information gain, comprising:

[0007] Construct training and testing datasets;

[0008] The object detection network is trained using the training dataset until convergence to obtain an initial unpruned teacher model, which is then tested using the test dataset.

[0009] Calculate the channel information gain values ​​of all BN layers to be pruned in the initial unpruned teacher model;

[0010] A preset pruning ratio is used to prune the network channel layers in the target detection network to obtain the student model.

[0011] The training loss function is modified, and the student model is fine-tuned using output knowledge distillation based on the modified training loss function to obtain the fine-tuned student model.

[0012] The fine-tuned student model was tested using the test dataset.

[0013] In this scheme, the target detection network that is trained on the training dataset until convergence is the YOLOv5 network.

[0014] In this scheme, calculating the channel information gain values ​​of all BN layers to be pruned in the initial unpruned teacher model specifically includes:

[0015] The output of the initial unpruned teacher model is transformed using the softmax activation function to obtain the distribution p(y|x,Θ);

[0016] Calculate the channel entropy H[p(y|x,Θ)] of the network layer to be pruned based on the distribution p(y|x,Θ);

[0017] Calculate the channel information gain value of the network layer to be pruned based on the calculated channel entropy H[p(y|x,Θ)].

[0018] The specific method for calculating the channel entropy of the network layer to be pruned is as follows:

[0019]

[0020] The specific calculation method for the channel information gain value of the network layer to be pruned is as follows:

[0021]

[0022] Where {x,y}∈{X,Y} and {X,Y} is the training dataset, R1(x,y) is the negligible first-order remainder, Θ is the initial unpruned teacher model, γ and β are the scaling factor and offset factor of the BN layer channel, respectively, k represents the k-th BN layer to be pruned, and i represents the i-th channel.

[0023] In this scheme, the preset pruning ratio is used to prune the network channel layers in the target detection network according to the preset pruning ratio to obtain the student model, specifically including:

[0024] Obtain the channel information gain value of the network layer to be pruned, and sort the channel information gain values ​​of all network layers to be pruned in descending order;

[0025] The pruning threshold is obtained based on the preset pruning ratio. The pruning threshold is compared with the information gain value of each channel of the network layer to be pruned, and the sparse mask M is used to mark whether the channel is pruned.

[0026] The student model is obtained by setting the scaling factor and offset factor of the BN channel of ratio p to 0 and deleting the relevant parameters of the channel to be pruned by M*Θ.

[0027] The specific rules for setting the mask M value are as follows:

[0028]

[0029] Where 0 represents trimmed channels, 1 represents preserved channels, and p represents the number of trimmed channels specified.

[0030] In this scheme, modifying the training loss function and fine-tuning the student model using output knowledge distillation based on the modified training loss function specifically includes:

[0031] The modified training loss function is as follows:

[0032] L final =L Yolo +λ D L Distillation ;

[0033] Among them, L Yolo Let λ be the loss function of the original YOLOv5 object detection network. D L represents the coefficient of the distillation loss function. Distillation This is the distillation loss function.

[0034] In this scheme, the loss function of the original YOLOv5 object detection network consists of the bounding box object confidence loss, classification loss, and coordinate regression loss, as specifically expressed below:

[0035]

[0036] in, and These represent the target confidence, class probability, and coordinate information contained in the predicted bounding box of each grid cell in the last layer feature map. and These represent the target confidence level, class probability, and coordinate information contained in the ground truth bounding box, respectively, f obj () represents the target confidence loss, f cl () represents the classification loss, f bb () represents the coordinate regression loss.

[0037] In this scheme, the distillation loss function consists of target confidence distillation loss, classification distillation loss, and coordinate regression distillation loss, specifically expressed as follows:

[0038]

[0039] in, These represent the target confidence, class probability, and coordinate information contained in the bounding box predicted by the teacher model.

[0040] A second aspect of the present invention also provides a computer-readable storage medium, characterized in that the computer-readable storage medium includes a target detection network channel pruning method program based on information gain, wherein when the program is executed by a processor, it implements the steps of the target detection network channel pruning method based on information gain as described in any of the preceding claims.

[0041] This invention discloses a channel pruning method for object detection networks based on information gain, comprising the following steps: constructing a training dataset and a test dataset; training the object detection network on the training dataset until convergence to obtain an initial unpruned teacher model; calculating the channel information gain values ​​of all BN layers to be pruned in the teacher model; presetting a pruning ratio, and pruning the network channel layers in the object detection network according to the preset pruning ratio to obtain a student model; modifying the training loss function, and using the teacher model to perform output knowledge distillation on the student model to achieve fine-tuning. This method does not require sparse training, the pruning process is simple and easy to implement, and achieves the effect of effectively compressing the object detection network under the premise of low loss accuracy. Attached Figure Description

[0042] To more clearly illustrate the technical solutions in the embodiments or examples of the present invention, the drawings used in the embodiments or examples will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained according to these drawings without creative effort.

[0043] Figure 1 A flowchart of a target detection network channel pruning method based on information gain is provided in an embodiment of the present invention;

[0044] Figure 2This is a comparison diagram of pruning results provided in an embodiment of the present invention;

[0045] Figure 3 This is a comparison chart of distillation results provided in an embodiment of the present invention;

[0046] The realization of the objective, functional features and advantages of the present invention will be further explained in conjunction with the embodiments and with reference to the accompanying drawings. Detailed Implementation

[0047] To better understand the above-mentioned objectives, features, and advantages of the present invention, the present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments. It should be noted that, unless otherwise specified, the embodiments and features described in these embodiments can be combined with each other.

[0048] Many specific details are set forth in the following description in order to provide a full understanding of the invention. However, the invention may also be practiced in other ways different from those described herein, and therefore the scope of protection of the invention is not limited to the specific embodiments disclosed below.

[0049] Figure 1 A flowchart of a target detection network channel pruning method based on information gain is provided in an embodiment of the present invention;

[0050] like Figure 1 As shown, the present invention provides a flowchart of a target detection network channel pruning method based on information gain, comprising:

[0051] S102, Construct the training dataset and the test dataset;

[0052] S104, The object detection network is trained using the training dataset until convergence is obtained to obtain an initial unpruned teacher model, which is then tested using the test dataset;

[0053] S106, Calculate the channel information gain values ​​of all BN layers to be pruned in the initial unpruned teacher model;

[0054] S108, preset pruning ratio, prune the network channel layer in the target detection network according to the preset pruning ratio to obtain the student model;

[0055] S110, Modify the training loss function, and fine-tune the student model using output knowledge distillation based on the modified training loss function to obtain the fine-tuned student model;

[0056] S112, The fine-tuned student model is tested using the test dataset.

[0057] It should be noted that this invention provides a channel pruning method for target detection networks based on information gain. First, the information gain value of the network layer channel to be pruned is calculated using the scaling factor and offset factor of the Batch Normalization (BN) layer in the target detection network. Then, the channel information gain values ​​are sorted in descending order, and a pruning threshold is obtained using a pre-set pruning ratio. Next, the network layer channels are marked and pruned. Finally, knowledge distillation is performed on the pruned network output. This method is simple and easy to implement, effectively reducing the number of parameters and computational load of the target detection network while minimizing the loss of detection accuracy.

[0058] Furthermore, in a preferred embodiment of the present invention, the training dataset and the test dataset are constructed using the PASCALVOC 2007 and 2012 datasets. The datasets are divided into four main categories: vehicle, household, animal, and person, with a total of 20 subcategories. The training set contains 16,551 images, and the test set contains 4,952 images.

[0059] Furthermore, in a preferred embodiment of the present invention, the target detection network that is trained on the target detection network using the training dataset until convergence is a YOLOv5 network.

[0060] Furthermore, in a preferred embodiment of the present invention, calculating the channel information gain values ​​of all BN layers to be pruned in the initial unpruned teacher model specifically includes:

[0061] The output of the initial unpruned teacher model is transformed using the softmax activation function to obtain the distribution p(y|x,Θ);

[0062] Calculate the channel entropy H[p(y|x,Θ)] of the network layer to be pruned based on the distribution p(y|x,Θ);

[0063] Calculate the channel information gain value of the network layer to be pruned based on the calculated channel entropy H[p(y|x,Θ)].

[0064] The specific method for calculating the channel entropy of the network layer to be pruned is as follows:

[0065]

[0066] The specific calculation method for the channel information gain value of the network layer to be pruned is as follows:

[0067]

[0068] Where {x,y}∈{X,Y} and {X,Y} is the training dataset, R1(x,y) is the negligible first-order remainder, Θ is the initial unpruned teacher model, γ and β are the scaling factor and offset factor of the BN layer channel, respectively, k represents the k-th BN layer to be pruned, and i represents the i-th channel.

[0069] Furthermore, in a preferred embodiment of the present invention, the preset pruning ratio, which involves pruning the network channel layers in the target detection network according to a preset pruning ratio to obtain the student model, specifically includes:

[0070] Obtain the channel information gain value of the network layer to be pruned, and sort the channel information gain values ​​of all network layers to be pruned in descending order;

[0071] The pruning threshold is obtained based on the preset pruning ratio. The pruning threshold is compared with the information gain value of each channel of the network layer to be pruned, and the sparse mask M is used to mark whether the channel is pruned.

[0072] The student model is obtained by setting the scaling factor and offset factor of the BN channel of ratio p to 0 and deleting the relevant parameters of the channel to be pruned by M*Θ.

[0073] The specific rules for setting the mask M value are as follows:

[0074]

[0075] Where 0 represents trimmed channels, 1 represents preserved channels, and p represents the number of trimmed channels specified.

[0076] Furthermore, in a preferred embodiment of the present invention, the modification of the training loss function, and the fine-tuning of the student model using output knowledge distillation based on the modified training loss function, specifically includes:

[0077] The modified training loss function is as follows:

[0078] L final =L Yolo +λ D L Distillation ;

[0079] Among them, L Yolo Let λ be the loss function of the original YOLOv5 object detection network. D L represents the coefficient of the distillation loss function. Distillation This is the distillation loss function.

[0080] It should be noted that λ D The coefficients of the distillation loss function are used to balance the original target detection loss and the distillation loss.

[0081] Furthermore, in a preferred embodiment of the present invention, the loss function of the original YOLOv5 object detection network consists of the bounding box object confidence loss, classification loss, and coordinate regression loss, specifically expressed as follows:

[0082]

[0083] in, and These represent the target confidence, class probability, and coordinate information contained in the predicted bounding box of each grid cell in the last layer feature map. and These represent the target confidence level, class probability, and coordinate information contained in the ground truth bounding box, respectively, f obj () represents the target confidence loss, f cl () represents the classification loss, f bb () represents the coordinate regression loss.

[0084] It should be noted that the target confidence loss f obj () and classification loss f cl () is achieved through the BCEWithLogitsLoss function, where the coordinate regression loss f is... bb () Achieved through the CIoU loss function

[0085] Furthermore, in a preferred embodiment of the present invention, the distillation loss function consists of target confidence distillation loss, classification distillation loss, and coordinate regression distillation loss, specifically represented as follows:

[0086]

[0087] in, These represent the target confidence, class probability, and coordinate information contained in the bounding box predicted by the teacher model.

[0088] It should be noted that the target confidence distillation loss function, the classification probability distillation loss function, and the bounding box coordinate regression distillation loss function all use the L2 loss function.

[0089] Figure 2 This is a comparison diagram of pruning results provided in an embodiment of the present invention. Figure 2 (a) It can be seen that the number of parameters (Parm), computational cost (FLOPs), and weight file size (Weights) of the pruned YOLOv5s-pruned network are 3.2M, 8.6GFLOPs, and 6.48M, respectively, which are compressed by 54.9%, 46.6%, and 53% compared to the original YOLOv5s network. Combined with... Figure 2(b) It can be seen that even after retraining and fine-tuning, the detection accuracy of the pruned network still suffers a high loss, with a loss of 0.05 mAP0.5.

[0090] Figure 3 This is a comparison chart of distillation results provided in an embodiment of the present invention. After knowledge distillation of the pruned network, YOLOv5s-kd is obtained, and its detection accuracy (mAP0.5) and mAP0.5:0.95 are both higher than those obtained through retraining and fine-tuning. Combined with... Figure 2 (b) It can be seen that the mAP0.5 after distillation is only 0.036 lower than the mAP0.5 before pruning. Combined with... Figure 2 (a) It can be seen that the target detection network channel pruning method based on information gain proposed in this invention can effectively compress the model with low loss accuracy. Thus, a large amount of background information output by the teacher model will affect the bounding box coordinate regression training problem of the student model.

[0091] A second aspect of the present invention also provides a computer-readable storage medium, characterized in that the computer-readable storage medium includes a target detection network channel pruning method program based on information gain, wherein when the program is executed by a processor, it implements the steps of the target detection network channel pruning method based on information gain as described in any of the preceding claims.

[0092] In the several embodiments provided in this application, it should be understood that the disclosed devices and methods can be implemented in other ways. The device embodiments described above are merely illustrative. For example, the division of units is only a logical functional division, and in actual implementation, there may be other division methods, such as: multiple units or components can be combined, or integrated into another system, or some features can be ignored or not executed. In addition, the coupling, direct coupling, or communication connection between the various components shown or discussed can be through some interfaces, and the indirect coupling or communication connection between devices or units can be electrical, mechanical, or other forms.

[0093] The units described above as separate components may or may not be physically separate. The components shown as units may or may not be physical units. They may be located in one place or distributed across multiple network units. Some or all of the units may be selected to achieve the purpose of this embodiment according to actual needs.

[0094] In addition, in the various embodiments of the present invention, each functional unit can be integrated into one processing unit, or each unit can be a separate unit, or two or more units can be integrated into one unit; the integrated unit can be implemented in hardware or in the form of hardware plus software functional units.

[0095] Those skilled in the art will understand that all or part of the steps of the above method embodiments can be implemented by hardware related to program instructions. The aforementioned program can be stored in a computer-readable storage medium. When the program is executed, it performs the steps of the above method embodiments. The aforementioned storage medium includes various media capable of storing program code, such as mobile storage devices, read-only memory (ROM), random access memory (RAM), magnetic disks, or optical disks.

[0096] Alternatively, if the integrated units of this invention are implemented as software functional modules and sold or used as independent products, they can also be stored in a computer-readable storage medium. Based on this understanding, the technical solutions of the embodiments of this invention, or the parts that contribute to the prior art, can be embodied in the form of a software product. This computer software product is stored in a storage medium and includes several instructions to cause a computer device (which may be a personal computer, server, or network device, etc.) to execute all or part of the methods described in the various embodiments of this invention. The aforementioned storage medium includes various media capable of storing program code, such as mobile storage devices, ROM, RAM, magnetic disks, or optical disks.

[0097] The above description is merely a specific embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any variations or substitutions that can be easily conceived by those skilled in the art within the technical scope disclosed in the present invention should be included within the scope of protection of the present invention. Therefore, the scope of protection of the present invention should be determined by the scope of the claims.

Claims

1. A channel pruning method for target detection networks based on information gain, characterized in that, include: Construct training and testing datasets; The object detection network is trained using the training dataset until convergence to obtain an initial unpruned teacher model, which is then tested using the test dataset. Calculate the channel information gain values ​​of all BN layers to be pruned in the initial unpruned teacher model; A preset pruning ratio is used to prune the network channel layers in the target detection network to obtain the student model. The training loss function is modified, and the student model is fine-tuned using output knowledge distillation based on the modified training loss function to obtain the fine-tuned student model. The fine-tuned student model was tested using the test dataset.

2. The target detection network channel pruning method based on information gain according to claim 1, characterized in that, The target detection network that is trained on the training dataset until convergence is the YOLOv5 network.

3. The target detection network channel pruning method based on information gain according to claim 1, characterized in that, The calculation of the channel information gain values ​​of all BN layers to be pruned in the initial unpruned teacher model specifically includes: The output of the initial unpruned teacher model is transformed using the softmax activation function to obtain the distribution p(y|x,Θ); Calculate the channel entropy H[p(y|x,Θ)] of the network layer to be pruned based on the distribution p(y|x,Θ); Calculate the channel information gain value of the network layer to be pruned based on the calculated channel entropy H[p(y|x,Θ)]. The specific method for calculating the channel entropy of the network layer to be pruned is as follows: The specific calculation method for the channel information gain value of the network layer to be pruned is as follows: Where {x,y}∈{X,Y} and {X,Y} is the training dataset, R1(x,y) is the negligible first-order remainder, Θ is the initial unpruned teacher model, γ and β are the scaling factor and offset factor of the BN layer channel, respectively, k represents the k-th BN layer to be pruned, and i represents the i-th channel.

4. The target detection network channel pruning method based on information gain according to claim 1, characterized in that, The preset pruning ratio is used to prune the network channel layers in the target detection network according to the preset pruning ratio to obtain the student model, specifically including: Obtain the channel information gain value of the network layer to be pruned, and sort the channel information gain values ​​of all network layers to be pruned in descending order; The pruning threshold is obtained based on the preset pruning ratio. The pruning threshold is compared with the information gain value of each channel of the network layer to be pruned, and the sparse mask M is used to mark whether the channel is pruned. The student model is obtained by setting the scaling factor and offset factor of the BN channel of ratio p to 0 and deleting the relevant parameters of the channel to be pruned by M*Θ. The specific rules for setting the mask M value are as follows: Where 0 represents trimmed channels, 1 represents preserved channels, and p represents the number of trimmed channels specified.

5. The target detection network channel pruning method based on information gain according to claim 1, characterized in that, The modification of the training loss function, and the fine-tuning of the student model using output knowledge distillation based on the modified training loss function, specifically includes: The modified training loss function is as follows: L final =L Yolo +λ D L Distillation ; Among them, L Yolo Let λ be the loss function of the original YOLOv5 object detection network. D L represents the coefficient of the distillation loss function. Distillation This is the distillation loss function.

6. The target detection network channel pruning method based on information gain according to claim 5, characterized in that, The loss function of the original YOLOv5 object detection network consists of the bounding box object confidence loss, classification loss, and coordinate regression loss, as specifically expressed below: in, and These represent the target confidence, class probability, and coordinate information contained in the predicted bounding box of each grid cell in the last layer feature map. and These represent the target confidence level, class probability, and coordinate information contained in the ground truth bounding box, respectively, f obj () represents the target confidence loss, f cl () represents the classification loss, f bb () represents the coordinate regression loss.

7. The target detection network channel pruning method based on information gain according to claim 5, characterized in that, The distillation loss function consists of target confidence distillation loss, classification distillation loss, and coordinate regression distillation loss, and is specifically expressed as follows: in, These represent the target confidence, class probability, and coordinate information contained in the bounding box predicted by the teacher model.

8. A computer-readable storage medium, characterized in that, The computer-readable storage medium includes a target detection network channel pruning method program based on information gain, which, when executed by a processor, implements the steps of the target detection network channel pruning method based on information gain as described in any one of claims 1 to 7.

Images