A real-time target detection method for UAV images based on channel pruning

Abstract

The present invention discloses a method for real-time target detection of unmanned aerial vehicle images based on channel pruning, using a pre-trained network as the initial network, using an updated loss function for sparse training, arranging the scale scaling factors of all batch normalization layers in order, and pruning the channels according to the sparse threshold; the channel pruning process uses a mask to mark the convolution layer channel, the channel mask to be pruned is 1, and the reserved channel mask is 0; the network layer is pruned layer by layer, and it is determined whether to delete the input, output, convolution kernel and batch normalization layer parameters connected to the channel according to the mask, and a new model parameter file is generated after the operation of the channel to be pruned is completed; finally, a smaller learning rate is used to fine-tune the model after channel pruning to restore the target recognition accuracy of the model. The present invention has low requirements on hardware resources and fast recognition speed, and can recognize the scene where the unmanned aerial vehicle is located in real time.

Description

Technical Field

[0001] The invention belongs to the technical field of unmanned aerial vehicle image target recognition, and in particular relates to a real-time target detection method for unmanned aerial vehicle images based on channel pruning. Background Art

[0002] When drones perform outdoor flight missions such as security protection, assembly monitoring, and natural exploration, they need to cooperate with ground stations to identify targets in real time for monitoring. When deploying to specific hardware resources, the hardware resources must be taken into consideration. The recognition speed must be further improved to achieve real-time recognition on laptop devices with limited performance. The computing power of drone ground stations is relatively redundant compared to embedded devices such as mobile phones, so we can compress the improved YOLO network model to reduce the number of parameters and improve the recognition speed. Channel pruning can achieve better generalization capabilities for various network models without relying on special computing resources. It can directly operate on convolutional layers and fully connected layers. Summary of the invention

[0003] Purpose of the invention: In order to overcome the problems of existing target recognition models that have high requirements for hardware resources, slow recognition speed, and difficulty in real-time recognition of the scene in which the drone is located, the present invention provides a real-time target detection method for drone images based on channel pruning.

[0004] Technical solution: The real-time target detection method for UAV images based on channel pruning described in the present invention comprises the following steps:

[0005] (1) Use the improved YOLO network to perform basic training on the drone dataset, and re-sparsely train the network after basic training based on the scaling factor of the batch normalization layer to generate a sparse scaling factor;

[0006] (2) In order to match the input and output feature channels of the residual module, a conservative pruning strategy and a full network pruning strategy are adopted, and channel pruning is performed using the scale factor of the BN layer as the selection criterion for the pruning channel;

[0007] (3) After pruning, the knowledge distillation strategy is used to fine-tune the pruned network and model to restore the model's target recognition accuracy;

[0008] (4) The model is comprehensively analyzed from two dimensions: model compression effect and target recognition effect, and the optimal implementation model for real-time multi-target recognition of UAV images is obtained.

[0009] Furthermore, the step (1) comprises the following steps:

[0010] (11) Introduce a scaling factor γ for each channel and multiply the output of the channel by the scaling factor;

[0011] (12) Jointly train the improved YOLO network weights and scaling factors, and perform sparse regularization on the scaling factors: The loss function of the channel pruning method based on the γ coefficient of the BN layer of the YOLO algorithm is as follows:

[0012] L bng =∑ (x，y) l(f(x,W),y)+λ∑ γ∈τ g(γ) (1)

[0013] Where (x, y) represents the input and target of the training, W represents the weight used for training, ∑ (x，y) l(f(x, W), y) is the loss value of the normal training of the convolutional neural network, the g function is the sparsity penalty term of the scaling factor, and λ is the coefficient that balances these two terms.

[0014] Furthermore, the implementation process of step (2) is as follows:

[0015] (21) Perform conservative pruning on the residual blocks with direct connection operations, that is, do not perform channel pruning to avoid inconsistent dimensions of the direct connection layers;

[0016] (22) Perform channel pruning on the general feature map, and finally prune the feature map associated with the residual block, that is, perform full network pruning;

[0017] (23) To perform channel pruning on directly connected feature tensors, the γ factors of channels at the same position need to be added and then sorted;

[0018] (24) Prune the channels of this feature map based on the scaling factor threshold.

[0019] Furthermore, the channel pruning described in step (2) uses a mask to mark the convolution layer channel, the channel mask to be pruned is 1, and the channel mask to be retained is 0; the network layer is pruned layer by layer, and it is determined whether to delete the input, output, convolution kernel and batch normalization layer parameters connected to the channel based on the mask, and a new model parameter file is generated after the operation of the channel to be pruned is completed.

[0020] Furthermore, the step (3) is implemented by the following formula:

[0021]

[0022] Among them, p represents the probability distribution of the true label, z and r represent the predicted outputs of the student network and the teacher network, and T is the temperature hyperparameter to make the output of the softmax classifier smoother and extract the knowledge of the label distribution from the output of the teacher network.

[0023] Beneficial effects: Compared with the prior art, the beneficial effects of the present invention are as follows: 1. Under the condition that the accuracy does not decrease after fine-tuning, the model compression effect of the full network channel pruning with a ratio of 0.35 is the best. After pruning, the number of parameters is reduced by 29.28 million, the amount of calculation is reduced by 26.4BFLOPs, the size of the target recognition model is reduced by 111.74M, the network running memory is reduced by 0.67G, and the network forward reasoning time is reduced by 3ms; 2. The number of parameters, the amount of calculation, the model memory size, and the forward reasoning time of the best pruned model are all less than those of conservative channel pruning. In the training environment of a desktop computer, the recognition speed can reach 33FPS, and in the environment of a laptop computer used in a ground station, the recognition speed can reach 25FPS. BRIEF DESCRIPTION OF THE DRAWINGS

[0024] Figure 1 is a flow chart of the present invention;

[0025] Figure 2 This is the channel pruning effect diagram based on the BN layer scaling factor;

[0026] Figure 3 The experimental results of the number distribution of scale factors in all BN layers under different balance factors in sparse training;

[0027] Figure 4 The schematic diagram of channel pruning. DETAILED DESCRIPTION

[0028] The present invention will be further described in detail below in conjunction with the accompanying drawings.

[0029] The present invention provides a real-time target detection method for UAV images based on channel pruning. Figure 1 As shown, the specific steps include:

[0030] Step 1: Use the improved YOLO network to perform basic training on the drone dataset, and re-sparsely train the network based on the scaling factor of the batch normalization layer to generate a sparse scaling factor.

[0031] like Figure 2 As shown in Figure 1, we first introduce a scaling factor γ for each channel and multiply the output of the channel by the scaling factor. Then we train the network weights and scaling factors together and perform sparse regularization on the scaling factors. Specifically, the loss function of the channel pruning method based on the γ coefficient of the BN layer can be expressed as follows:

[0032] L bng =∑ (x，y) l(f(x,W),y)+λ∑ γ∈τ g(γ) (1)

[0033] Among them, (x, y) represents the input and target of the training, W represents the weight used for training, and the previous summation corresponds to the loss value of the normal training of the convolutional neural network; the g function is the sparsity penalty term of the scaling factor, and λ is the coefficient that balances these two terms. You can choose the L1 regularization term or the L2 regularization term as the penalty term of the scaling factor. These two regularization methods are widely used to achieve sparsity, where the L1 regularization term is g(γ)=|γ| and the L2 regularization term is g(γ)=γ 2 Taking the L1 regularization term as an example, when using a non-smooth L1 penalty term, sub-gradient descent can be used as an optimization method. If a smooth L1 penalty term is used, sub-gradient descent can be avoided at non-smooth points. The network and model with improved structure are trained sparsely. During sparse training, the number of scale factors of all BN layers is counted, which reflects the distribution change of scale factors. The experimental results are shown in Figure 3 shown.

[0034] Step 2: In order to match the input and output feature channels of the residual module, a conservative pruning strategy and a full network pruning strategy are adopted, and channel pruning is performed using the scale scaling factor of the BN layer as the selection criterion for the pruning channel.

[0035] like Figure 4 As shown in the figure, the YOLO network structure can be divided into the backbone network and the detection part. When performing channel pruning on the backbone part, the residual module is a structure that requires special attention. In order to solve the problem of gradient divergence, the residual module introduces a direct connection operation, which adds two feature tensors of the same dimension bit by bit. If channel pruning is performed directly, the input and output feature maps of the residual module cannot match the number of channels. Therefore, the channels cannot be directly deleted for these feature tensors, and the channels in the same position must be left. The first strategy is to not perform channel pruning on the residual blocks with direct connection operations directly, thereby avoiding the problem of inconsistent dimensions of the direct connection layer, which is conservative pruning. The second strategy is to perform channel pruning on the general feature map first, and then prune the feature map associated with the residual block, which is called full network pruning. To perform channel pruning on the directly connected feature tensor, it is necessary to add and sort the γ factors of the channels in the same position, and finally prune the channels of this feature map based on the scaling factor threshold.

[0036] During the channel pruning process, masks are used to mark the convolution layer channels. The mask of the channel that needs to be pruned is 1, and the mask of the retained channel is 0. The network layer is pruned layer by layer, and the mask is used to determine whether to delete the input, output, convolution kernel, and batch normalization layer parameters connected to the channel. After the pruned channel operation is completed, a new model parameter file is generated.

[0037] Step 3: After pruning, the knowledge distillation strategy is used to fine-tune the pruned network and model to restore the target recognition accuracy of the model.

[0038] Using the knowledge distillation strategy, knowledge distillation is suitable for networks with similar model size and structure, and has significant effects in fine-tuning after pruning. The knowledge distillation strategy uses the feature map of the teacher network to train a compact student network. In the pruning and fine-tuning stage, the teacher network is a pre-trained model before the pruning operation. The pruned network continuously improves the target recognition accuracy by imitating the pre-trained model while keeping the complexity of the pruned model unchanged. The specific implementation method is to increase the distillation loss during training to penalize the inconsistency of the softmax classifier outputs of the two networks. Originally, the negative cross entropy loss l(p, softmax(z)) was used to measure the difference between the predicted output of the network and the true label. Now the loss function of the distillation part is added, so the loss function of the knowledge distillation becomes as shown in the following formula (2):

[0039]

[0040] Where p represents the probability distribution of the true label, z and r represent the predicted outputs of the student network and the teacher network, and T is the temperature hyperparameter to make the output of the softmax classifier smoother and extract the knowledge of the label distribution from the output of the teacher network.

[0041] Step 4: Comprehensively analyze the model from two dimensions: model compression effect and target recognition effect, and obtain the optimal implementation model for real-time multi-target recognition of drone images. The full network pruning strategy used is to find the mask information of the convolution layer with the global threshold as the indicator. For each group of direct connection operations, the pruning masks of the connected convolution layers are combined, and the fused mask is used to determine the pruning. This method constrains the reserved channels of each layer. Adding processing to the activation offset value during execution can reduce the accuracy loss caused by pruning. Pruning is performed at a ratio of 0.3, 0.35, 0.4, 0.45, 0.5, and 0.55, respectively. The model compression index and target recognition performance index of the pruned network and model are shown in Table 1 and Table 2, respectively.

[0042] Table 1 Experimental results of model compression indicators of full network pruning strategy

[0043]

[0044] Table 2 Experimental results of target recognition indicators for the full network pruning strategy

[0045]

[0046] From Tables 1 and 2, we can see that under the condition that the accuracy does not decrease after fine-tuning, the model compression effect of full network channel pruning with a ratio of 0.35 is the best. After pruning, the number of parameters is reduced by 29.28 million, the amount of calculation is reduced by 26.4BFLOPs, the size of the target recognition model is reduced by 111.74M, the network running memory is reduced by 0.67G, and the network forward inference time is reduced by 3ms.

Claims

1. A real-time target detection method for UAV images based on channel pruning, It is characterized in that The following steps are involved: (1) Use the improved YOLO network to perform basic training on the drone dataset, and re-sparsely train the network based on the scaling factor of the batch normalization layer to generate a sparse scaling factor; (2) In order to match the input and output feature channels of the residual module, a conservative pruning strategy and a full network pruning strategy are adopted, and channel pruning is performed using the scale factor of the BN layer as the selection criterion for the pruning channel; (3) After pruning, the knowledge distillation strategy is used to fine-tune the pruned network and model to restore the model's target recognition accuracy; (4) Comprehensively analyze the model from two dimensions: model compression effect and target recognition effect, and obtain the optimal implementation model for real-time multi-target recognition of UAV images; The implementation process of step (2) is as follows: (21) Perform conservative pruning on the residual blocks with direct connection operations, that is, do not perform channel pruning to avoid inconsistent dimensions of the direct connection layers; (22) Perform channel pruning on the general feature map, and finally prune the feature map associated with the residual block, that is, perform full network pruning; (23) To perform channel pruning on directly connected feature tensors, the γ factors of channels at the same position need to be added and then sorted; (24) Pruning the channels of this feature map based on the scaling factor threshold; The step (3) is implemented by the following formula: Among them, p represents the probability distribution of the true label, z and r represent the predicted outputs of the student network and the teacher network, and T is the temperature hyperparameter to make the output of the softmax classifier smoother and extract the knowledge of the label distribution from the output of the teacher network.

2. The real-time target detection method for UAV images based on channel pruning according to claim 1, It is characterized in that The step (1) comprises the following steps: (11) Introduce a scaling factor γ for each channel and multiply the output of the channel by the scaling factor; (12) Jointly train the improved YOLO network weights and scaling factors, and perform sparse regularization on the scaling factors: The loss function of the channel pruning method based on the γ coefficient of the BN layer of the YOLO algorithm is as follows: L bng =∑ (x,y) l(f(x,W),y)+λ∑ γ∈τ g(γ) (1) Among them, (x, y) represents the input and target of the training, W represents the weight used for training, ∑ (x,y) l(f(x,W),y) is the loss value of the normal training of the convolutional neural network, the g function is the sparsity penalty term of the scaling factor, and λ is the coefficient that balances these two terms.

3. The real-time target detection method for UAV images based on channel pruning according to claim 1, It is characterized in that The channel pruning described in step (2) uses a mask to mark the convolution layer channel. The mask of the channel to be pruned is 1, and the mask of the retained channel is 0. The network layer is pruned layer by layer, and the input, output, convolution kernel and batch normalization layer parameters connected to the channel are determined according to the mask. After the operation of the channel to be pruned is completed, a new model parameter file is generated.

Images