Target detection network channel pruning method based on information gain

Abstract

The invention discloses a target detection network channel pruning method based on information gain. The method comprises the following steps: constructing a training data set and a test data set; training the target detection network on the training data set until convergence, and obtaining an initial unpruned teacher model; calculating channel information gain values of all BN layers to be pruned in the teacher model; presetting a pruning proportion, and pruning a network channel layer in the target detection network according to the preset pruning proportion to obtain a student model; the training loss function is modified, and the teacher model is utilized to perform output knowledge distillation on the student model to achieve the fine adjustment purpose. According to the method, sparse training is not needed, the pruning process is simple and easy to implement, and the effect of effectively compressing the target detection network on the premise of low loss precision is achieved.

Description

Technical Field

[0001] The present invention relates to the technical field of target detection, and in particular to a target detection network channel pruning method based on information gain. Background Art

[0002] Object detection is the task of classifying and locating objects in images or videos. It has broad application prospects and application value. Object detection can be directly applied to many practical scenarios such as intelligent transportation, industrial equipment detection, defect detection, etc. Taking the application of object detection in the field of transportation as an example, in terms of traffic monitoring and security, object detection algorithms can be used to analyze monitoring video streams and detect traffic violations on the road in real time, such as driving in the wrong direction, running red lights, and crossing the line, thereby improving monitoring efficiency and traffic safety; in terms of intelligent traffic control, object detection technology can be used to automatically perceive traffic flow and realize intelligent traffic control and scheduling, such as signal light control, priority passage, and traffic jam prediction.

[0003] Since deep neural networks have a large number of parameters, some of them that contribute little to the output of the network model are considered redundant parameters. These parameters can be pruned to reduce the amount of computation and parameters of the network model. Most of the current channel pruning solutions for target detection networks require sparse training first, and multiple fine-tuning of parameters is required to achieve a good sparse effect. Most channel pruning methods use traditional retraining methods to fine-tune the pruned network to restore accuracy. However, because the retraining and fine-tuning methods have limited ability to restore accuracy, knowledge distillation is used to fine-tune the pruned network. Knowledge distillation is a method of transferring knowledge from a pre-trained complex teacher model to a student model, so that a smaller student model can replace a larger teacher model during deployment.

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

[0005] The present invention overcomes the defects of the prior art and provides a target detection network channel pruning method based on information gain, the important purpose of which is to reduce the number of parameters and calculations of the target detection network with lower loss of detection accuracy.

[0006] To achieve the above-mentioned purpose, the first aspect of the present invention provides a target detection network channel pruning method based on information gain, comprising:

[0007] Construct training and testing datasets;

[0008] Using the training data set to train the target detection network until convergence, obtaining an initial unpruned teacher model, and testing it using the test data set;

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

[0010] Preset a pruning ratio, and prune the network channel layer in the target detection network according to the preset pruning ratio to obtain a student model;

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

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

[0013] In this solution, the target detection network that is trained using the training data set until convergence is a YOLOv5 network.

[0014] In this solution, the channel information gain values ​​of all the BN layers to be pruned in the initial unpruned teacher model are calculated, specifically including:

[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 calculation method of the channel entropy of the network layer to be pruned is as follows:

[0019]

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

[0021]

[0022] Among them, {x,y}∈{X,Y} and {X,Y} are training data sets, R1(x,y) is a first-order residue that can be ignored, Θ is the initial unpruned teacher model, γ and β are the scaling factor and offset factor of the BN layer channel respectively, k represents the kth BN layer to be pruned, and i represents the i-th channel.

[0023] In this solution, the preset pruning ratio is used to prune the network channel layer in the target detection network according to the preset pruning ratio to obtain a student model, which specifically includes:

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

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

[0026] The scaling factor and offset factor of the BN channel of ratio p are set to 0 through M*Θ and the relevant parameters of the channel to be pruned are deleted to obtain the student model;

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

[0028]

[0029] Among them, 0 represents pruned channels, 1 represents retained channels, and p represents the number of specified pruned channels.

[0030] In this solution, the training loss function is modified, and the student model is fine-tuned by using the output knowledge distillation method based on the modified training loss function, specifically including:

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

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

[0033] Among them, L Yolo is the loss function of the original YOLOv5 target detection network, λ D is the coefficient of the distillation loss function, L Distillation is the distillation loss function.

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

[0035]

[0036] in, and The target confidence, category probability and coordinate information contained in the bounding box predicted for each grid unit of the last layer feature map are respectively, and are the target confidence, category probability and coordinate information contained in the real annotation box, respectively, obj () is the target confidence loss, f cl () is the classification loss, f bb () is the coordinate regression loss.

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

[0038]

[0039] in, The target confidence, category probability and coordinate information contained in the bounding box predicted by the teacher model respectively.

[0040] The 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, and when the target detection network channel pruning method program based on information gain is executed by a processor, the steps of the target detection network channel pruning method based on information gain as described in any one of the above items are implemented.

[0041] The present invention discloses a target detection network channel pruning method based on information gain, comprising the following steps: constructing a training data set and a test data set; training the target detection network on the training data set 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 layer in the target 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 the purpose of fine-tuning. The method does not require sparse training, and the pruning process is simple and easy to implement, achieving the effect of effectively compressing the target detection network under the premise of low loss accuracy. BRIEF DESCRIPTION OF THE DRAWINGS

[0042] In order to more clearly illustrate the technical solutions in the embodiments or exemplary embodiments of the present invention, the drawings required for use in the embodiments or exemplary descriptions will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained according to the drawings without paying creative work.

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

[0044] Figure 2A comparison diagram of pruning results provided by an embodiment of the present invention;

[0045] Figure 3 A comparison chart of distillation results provided by an embodiment of the present invention;

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

[0047] In order to more clearly understand the above-mentioned purpose, features and advantages of the present invention, the present invention is further described in detail below in conjunction with the accompanying drawings and specific embodiments. It should be noted that the embodiments of the present application and the features in the embodiments can be combined with each other without conflict.

[0048] In the following description, many specific details are set forth to facilitate a full understanding of the present invention. However, the present invention may also be implemented in other ways different from those described herein. Therefore, the protection scope of the present invention is not limited to the specific embodiments disclosed below.

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

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

[0051] S102, constructing a training data set and a test data set;

[0052] S104, training the target detection network using the training data set until convergence, obtaining an initial unpruned teacher model, and testing it using the test data set;

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

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

[0055] S110, modifying the training loss function, and fine-tuning the student model by using an output knowledge distillation method based on the modified training loss function to obtain a fine-tuned student model;

[0056] S112, using the test data set to test the fine-tuned student model.

[0057] It should be noted that the present invention provides a target detection network channel pruning method based on information gain, firstly, the information gain value of the network layer channel to be pruned is calculated using the scaling factor and offset factor of the BN layer in the target detection network, then the channel information gain value is sorted in descending order and the pruning threshold is obtained using the pre-set pruning ratio, then the network layer channel is marked and pruned, and finally the output knowledge distillation is performed on the pruned network. The method is simple and easy to implement, and can effectively reduce the number of parameters and calculations of the target detection network and the detection accuracy with low loss.

[0058] Furthermore, in a preferred embodiment of the present invention, the training data set and the test data set are constructed using the PASCAL VOC 2007 and 2012 data sets, which are divided into four major 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 using the training data set until convergence is a YOLOv5 network.

[0060] Further, in a preferred embodiment of the present invention, the calculating of the channel information gain values ​​of all the 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 calculation method of the channel entropy of the network layer to be pruned is as follows:

[0065]

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

[0067]

[0068] Among them, {x,y}∈{X,Y} and {X,Y} are training data sets, R1(x,y) is a first-order residue that can be ignored, Θ is the initial unpruned teacher model, γ and β are the scaling factor and offset factor of the BN layer channel respectively, k represents the kth BN layer to be pruned, and i represents the i-th channel.

[0069] Further, in a preferred embodiment of the present invention, the preset pruning ratio, pruning the network channel layer in the target detection network according to the 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 the network layers to be pruned in descending order;

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

[0072] The scaling factor and offset factor of the BN channel of ratio p are set to 0 through M*Θ and the relevant parameters of the channel to be pruned are deleted to obtain the student model;

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

[0074]

[0075] Among them, 0 indicates pruned channels, 1 indicates retained channels, and p indicates the number of specified pruned channels.

[0076] Furthermore, in a preferred embodiment of the present invention, the training loss function is modified, and the student model is fine-tuned by using the output knowledge distillation method based on the modified training loss function, specifically including:

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

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

[0079] Among them, L Yolo is the loss function of the original YOLOv5 target detection network, λ D is the coefficient of the distillation loss function, L Distillation is the distillation loss function.

[0080] It should be noted that λ D is the coefficient of the distillation loss function, which is used to balance the original target detection loss and the distillation loss.

[0081] Further, in a preferred embodiment of the present invention, the loss function of the original YOLOv5 target detection network is composed of target confidence loss, classification loss and coordinate regression loss of the bounding box, which are specifically expressed as follows:

[0082]

[0083] in, and The target confidence, category probability and coordinate information contained in the bounding box predicted for each grid unit of the last layer feature map are respectively, and are the target confidence, category probability and coordinate information contained in the real annotation box, respectively, obj () is the target confidence loss, f cl () is the classification loss, f bb () is the coordinate regression loss.

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

[0085] Furthermore, in a preferred embodiment of the present invention, the distillation loss function is composed of target confidence distillation loss, classification distillation loss and coordinate regression distillation loss, which are specifically expressed as follows:

[0086]

[0087] in, The target confidence, category probability and coordinate information contained in the bounding box predicted by the teacher model respectively.

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

[0089] Figure 2 A comparison chart of pruning results provided by an embodiment of the present invention. Figure 2 (a) It can be seen that the parameters (Paramers, Parm), calculation amount (FLOPs) and weight file size (Weights) of the pruned YOLOv5s-pruned network are 3.2M, 8.6GFLOPs and 6.48M respectively, which are 54.9%, 46.6% and 53% smaller than those of the original YOLOv5s network. Figure 2(b) It can be seen that even after retraining and fine-tuning, the detection accuracy of the pruned network still suffers from a high loss, with a mAP0.5 loss of 0.05.

[0090] Figure 3 A comparison chart of distillation results provided by an embodiment of the present invention. After performing knowledge distillation on the pruned network, YOLOv5s-kd is obtained, and its detection accuracy mAP0.5 and mAP0.5:0.95 are both higher than those obtained by retraining and fine-tuning. Figure 2 (b) It can be seen that the mAP0.5 after distillation is only 0.036 lower than the mAP0.5 before pruning. Figure 2 (a) It can be seen that the information gain-based target detection network channel pruning method proposed in the present invention can effectively compress the model under the premise of low loss accuracy, so that 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] The 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, and when the target detection network channel pruning method program based on information gain is executed by a processor, the steps of the target detection network channel pruning method based on information gain as described in any one of the above items are implemented.

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

[0093] The units described above as separate components may or may not be physically separated, and the components displayed as units may or may not be physical units; they may be located in one place or distributed on multiple network units; some or all of the units may be selected according to actual needs to achieve the purpose of the present embodiment.

[0094] In addition, all functional units in the embodiments of the present invention may be integrated into one processing unit, or each unit may be separately used as a unit, or two or more units may be integrated into one unit; the above-mentioned integrated units may be implemented in the form of hardware or in the form of hardware plus software functional units.

[0095] Those skilled in the art can understand that: all or part of the steps of implementing the above method embodiments can be completed by hardware related to program instructions, and the aforementioned program can be stored in a computer-readable storage medium. When the program is executed, it executes the steps of the above method embodiments; and the aforementioned storage medium includes: mobile storage devices, read-only memories (ROM, Read-Only Memory), random access memories (RAM, Random Access Memory), disks or optical disks, and other media that can store program codes.

[0096] Alternatively, if the above-mentioned integrated unit of the present invention is implemented in the form of a software function module and sold or used as an independent product, it can also be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the embodiment of the present invention can be essentially or partly reflected in the form of a software product that contributes to the prior art. The computer software product is stored in a storage medium and includes several instructions for a computer device (which can be a personal computer, a server, or a network device, etc.) to execute all or part of the methods described in each embodiment of the present invention. The aforementioned storage medium includes: various media that can store program codes, such as mobile storage devices, ROM, RAM, magnetic disks or optical disks.

[0097] The above is only a specific embodiment of the present invention, but the protection scope of the present invention is not limited thereto. Any person skilled in the art can easily think of changes or substitutions within the technical scope disclosed by the present invention, which should be included in the protection scope of the present invention. Therefore, the protection scope of the present invention should be based on the protection scope of the claims.

Claims

1. A target detection network channel pruning method based on information gain, characterized in that: include: Construct training and testing datasets; Using the training data set to train the target detection network until convergence, obtaining an initial unpruned teacher model, and testing it using the test data set; Calculate the channel information gain values ​​of all BN layers to be pruned in the initial unpruned teacher model; Preset a pruning ratio, and prune the network channel layer in the target detection network according to the preset pruning ratio to obtain a student model; The training loss function is modified, and based on the modified training loss function, the student model is fine-tuned by using the output knowledge distillation method to obtain a fine-tuned student model; The fine-tuned student model is tested using the test dataset.

2. According to the information gain-based target detection network channel pruning method of claim 1, it is characterized in that: The target detection network that is trained using the training data set until convergence is a YOLOv5 network.

3. According to the information gain-based target detection network channel pruning method of claim 1, it is characterized in that: The calculating of the channel information gain values ​​of all the 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 calculation method of the channel entropy of the network layer to be pruned is as follows: The specific calculation method of the channel information gain value of the network layer to be pruned is as follows: Among them, {x,y}∈{X,Y} and {X,Y} are training data sets, R1(x,y) is a first-order residue that can be ignored, Θ is the initial unpruned teacher model, γ and β are the scaling factor and offset factor of the BN layer channel respectively, k represents the kth BN layer to be pruned, and i represents the i-th channel.

4. According to the information gain-based target detection network channel pruning method of claim 1, it is characterized in that: The preset pruning ratio, pruning the network channel layer in the target detection network according to the preset pruning ratio to obtain the student model, specifically includes: Obtain the channel information gain value of the network layer to be pruned, and sort the channel information gain values ​​of all the network layers to be pruned in descending order; A pruning threshold is obtained based on a preset pruning ratio, the pruning threshold is compared with the information gain value of each channel of the network layer to be pruned, and a sparse mask M is used to mark whether the channel is pruned; The scaling factor and offset factor of the BN channel of ratio p are set to 0 through M*Θ and the relevant parameters of the channel to be pruned are deleted to obtain the student model; The setting rules of the mask M value are as follows: Among them, 0 represents pruned channels, 1 represents retained channels, and p represents the number of specified pruned channels.

5. According to the information gain-based target detection network channel pruning method of claim 1, it is characterized in that: The training loss function is modified, and the student model is fine-tuned by using the output knowledge distillation method based on the modified training loss function, specifically including: The modified training loss function is as follows: L final =L Yolo +λ D L Distillation ; Among them, L Yolo is the loss function of the original YOLOv5 target detection network, λ D is the coefficient of the distillation loss function, L Distillation 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 target detection network consists of the target confidence loss of the bounding box, the classification loss, and the coordinate regression loss, which are specifically expressed as follows: in, and The target confidence, category probability and coordinate information contained in the bounding box predicted for each grid unit of the last layer feature map are respectively, and are the target confidence, category probability and coordinate information contained in the real annotation box, respectively, obj () is the target confidence loss, f cl () is the classification loss, f bb () is the coordinate regression loss.

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

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. When the target detection network channel pruning method program based on information gain is executed by a processor, the steps of the target detection network channel pruning method based on information gain as described in any one of claims 1 to 7 are implemented.

Images