Lightweight target detection method fusing attention mechanism

Abstract

The invention relates to a lightweight target detection method fusing an attention mechanism, and belongs to the technical field of image recognition. The method comprises the following steps: firstly, making a target image data set and preprocessing; secondly, combining an MBConv structure and a CA attention mechanism to construct a CSPCAMBConv structure, and using the CSPCAMBConv structure to reconstruct a backbone network of Yolov5s; a bidirectional feature pyramid fusion module is adopted in the neck network, and feature information transmission and sharing are enhanced; besides, a scale inconsistency item and a width difference item are introduced into a loss function to balance the difference between a prediction frame and a real frame, so that the model is more sensitive to the scale difference of the target in the training process. And finally, pruning the model through sparse training to compress the size of the network model so as to obtain an optimal lightweight model. The method is suitable for an edge terminal with limited computing resources, targets of different scales can be accurately positioned under a complex background, and the problems of wrong detection and missing detection of the targets are effectively reduced.

Description

Technical Field

[0001] The present invention belongs to the technical field of image recognition terminals and relates to a lightweight target detection method integrating an attention mechanism. Background Art

[0002] In recent years, deep neural networks based on convolution operations have achieved superior performance in many fields of machine vision and natural language processing due to their good generalization ability and efficient feature extraction ability. Machine vision target detection algorithms can quickly locate and classify specific targets, and target detection technology based on convolutional neural networks is widely used in military, industrial quality inspection and other fields. Although deep neural networks are widely used, there is still a lot of room for development in engineering deployment. Deep neural networks have high time and space complexity, consume a lot of computing and storage resources, and have high requirements for hardware computing resources and storage, making them difficult to deploy on terminals with limited resources. Therefore, in order to ensure that the target detection algorithm can run under resource-constrained conditions, the detection network must be lightweight.

[0003] In addition, as an important derivative of the Internet, the Internet of Things provides the ability of interconnection and intelligent transmission. Against the background of the continuous rise of a large number of mobile terminals, the target detection model is gradually miniaturized and shifted from traditional server deployment to edge terminal platform deployment. However, with limited computing resources on edge terminals, how to improve the speed and accuracy of target detection is still an urgent problem to be solved. Therefore, under the premise of ensuring the accuracy of target detection, reducing the number of parameters and calculations of the network model, and deploying the target recognition algorithm model on a hardware platform with limited resources for stable operation has practical application value and broad engineering application prospects. Summary of the invention

[0004] In view of this, the purpose of the present invention is to provide a lightweight target detection method integrating an attention mechanism to solve the problem of deploying a high-precision target detection model on a terminal with scarce computing resources.

[0005] In order to achieve the above object, the present invention provides the following technical solutions:

[0006] A lightweight object detection method integrating an attention mechanism, the method comprising:

[0007] S1. Obtain a public target image dataset and preprocess the dataset, including dataset augmentation, annotation, and partitioning, and finally select training and test datasets;

[0008] S2. Use the MBConv structure in the EfficientNet network, combine it with the CA attention mechanism to build the CSP_CAMBConv structure, improve the activation function in the CA attention mechanism, and use the CSP_CAMBConv structure to reconstruct the backbone network of Yolov5s;

[0009] S3, in the neck network of Yolov5s, a path aggregation structure is used to fuse feature information of different levels, and scale inconsistency terms and width difference terms are introduced into the CIoU loss function to construct the SCIoU loss function;

[0010] S4. Perform sparse training on the improved target detection model to facilitate subsequent model pruning, and use the global attenuation method to adjust the sparse factor during the sparse process to quickly achieve the sparseness of the sparse factor. Finally, fine-tune the network to obtain the optimal lightweight target detection model.

[0011] S5. Acquire the image and input it into the optimal target detection model for target detection.

[0012] Furthermore, in step S2, in the constructed CSP_CAMBConv structure, the input is first processed by a CBS module and then passed to the Concat module and the CAMBConv structure respectively. The feature map is spliced ​​by the Concat module, and the features are acquired by the CAMBConv structure and then passed to the Concat module. The spliced ​​feature map is further extracted with standardized features by another CBS module.

[0013] Furthermore, in step S2, the constructed CSP_CAMBConv is applied to the Yolov5s backbone network to improve Yolov5s. Specifically, CSP_CAMBConv is applied to the 4th, 6th and 8th CSP modules of the Yolov5s backbone network. In addition, the MBConv structure in the EfficientNet feature extraction network is applied to the ordinary convolution modules in the backbone network, namely the 2nd, 3rd, 5th, 7th and 9th layers, to form a new Yolov5s backbone network structure.

[0014] Furthermore, in step S2, improving the activation function in the CA attention mechanism includes changing the tan function in the original GELU activation function to the erf function to obtain an improved activation function:

[0015]

[0016] Where erf(x) represents the error function and x represents the input feature map.

[0017] Furthermore, in step S3, a path aggregation structure is used in the neck network of Yolov5s to fuse feature information of different levels. Specifically, the feature maps output by the 5th and 7th layers in the Yolov5s backbone network are respectively input into the 1st and 2nd feature fusion modules in the neck network, and the feature map output by the 10th layer of the backbone network is output to the first CBS module in the neck network after the SPPF module; the outputs of the 1st CBS module and the 2nd CBS module in the neck network are respectively used as the inputs of the 3rd and 4th feature fusion modules in the neck network.

[0018] Furthermore, in step S3, scale inconsistency terms and width difference terms are introduced into the CIoU loss function to construct the SCIoU loss function. Specifically, scale inconsistency terms and width difference terms are introduced on the basis of the CIoU loss function to weigh the difference between the predicted box and the true box, and make the model more sensitive to the scale difference of the target during training. The improved SCIoU loss function is as follows:

[0019]

[0020] In the formula, IoU represents the intersection over union ratio between the predicted box and the true box; b represents the center point of the predicted bounding box, and b gt represents the center point of the true bounding box; ρ represents the Euclidean distance; w gt and h gt Represent the width and height of the real bounding box respectively; w pred and h pred Represent the width and height of the predicted bounding box, respectively; c represents the diagonal distance of the minimum bounding rectangle, and ε represents a small constant to avoid zero division errors. The improved SCIoU loss function measures the scale difference of the target by logarithmically transforming the area ratio of the predicted box and the true box. The logarithmic operation makes the model more sensitive when the scale difference is large; and the width difference term can reduce the impact of the width difference on the overall loss function through the ratio of the difference between the maximum and minimum values.

[0021] Furthermore, in step S4, the improved Yolov5s is initially trained to obtain a basic model; then the initial model is sparsely trained to highlight important features; then the local structure of the model is pruned to remove unimportant connections; and then the pruned network model is fine-tuned and optimized to obtain an optimized pruned model.

[0022] At the same time, in the sparse training, global attenuation is introduced to dynamically adjust the sparse factor, so that the regularization weight gradually decreases with the training process, and the sparsification of the parameter γ is quickly achieved; the improved sparse training loss function is shown as follows:

[0023]

[0024] In the formula, λ represents the sparse coefficient corresponding to the pruning sparsification of the model; l represents the training loss function before the model introduces regularization; x′ represents the training input data; f represents the model forward prediction function; y represents the true label; t represents the current training round number; T represents the total number of training rounds; Γ represents the scaling factor of the BN layer in the model. The basic linear attenuation term is It is the core attenuation part. As the number of training times increases, the sparse coefficient λ will gradually decrease. Dynamic fluctuation term Controlled by β, greater regularization can be applied at the beginning of training, exerting a greater influence on the sparse coefficient λ in the early and middle stages of training, while gradually reducing the influence in the later stages of training to adapt to the overall attenuation trend.

[0025] The beneficial effects of the present invention are:

[0026] (1) The present invention adopts the MBConv structure in the EfficientNet network and combines it with the improved CA attention mechanism to construct the CSP_CAMBConv structure, and uses this structure to reconstruct the backbone network of Yolov5s. This improvement not only takes into account the relationship between space and channels, but also reduces the number of parameters and computational complexity of the model.

[0027] (2) To meet the detection requirements of targets of different scales, the present invention adopts a bidirectional feature pyramid fusion module in the neck network to enhance the transmission and sharing of feature information. In addition, a scale inconsistency term and a width difference term are introduced into the loss function to weigh the difference between the predicted box and the true box, and make the model more sensitive to the scale difference of the target during training.

[0028] (3) The present invention is applicable to edge terminals with limited computing resources and can accurately locate targets of different scales in complex backgrounds, effectively reducing the problems of false detection and missed detection of targets.

[0029] Other advantages, objectives and features of the present invention will be described in the following description to some extent, and to some extent, will be obvious to those skilled in the art based on the following examination and study, or can be taught from the practice of the present invention. The objectives and other advantages of the present invention can be realized and obtained through the following description. BRIEF DESCRIPTION OF THE DRAWINGS

[0030] In order to make the purpose, technical solutions and advantages of the present invention more clear, the present invention will be described in detail below in conjunction with the accompanying drawings, wherein:

[0031] Figure 1 It is a schematic diagram of the process of the method of the present invention;

[0032] Figure 2A schematic diagram of the CSP_CAMBConv structure proposed in one embodiment of the present invention;

[0033] Figure 3 A diagram of a lightweight network structure proposed in one embodiment of the present invention;

[0034] Figure 4 Flowchart for model pruning. DETAILED DESCRIPTION

[0035] The following describes the embodiments of the present invention by specific examples, and those skilled in the art can easily understand other advantages and effects of the present invention from the contents disclosed in this specification. The present invention can also be implemented or applied through other different specific embodiments, and the details in this specification can also be modified or changed in various ways based on different viewpoints and applications without departing from the spirit of the present invention. It should be noted that the illustrations provided in the following embodiments only illustrate the basic concept of the present invention in a schematic manner, and the following embodiments and the features in the embodiments can be combined with each other without conflict.

[0036] like Figure 1 As shown, a lightweight target detection method integrating an attention mechanism is provided in one embodiment of the present invention, and the specific contents of the method are as follows:

[0037] S1. Obtain a public target image dataset and preprocess the dataset, including dataset augmentation, labeling, and partitioning, and finally select the training and test datasets.

[0038] In this example, the public dataset VisDrone2019 is used as the test, training, and validation dataset. A variety of data enhancement preprocessing methods are used to generate more diverse samples, including amplifying the dataset by rotation, mirroring, noise, light and dark changes, etc., and filtering out invalid images that do not contain the target in the picture due to preprocessing, and finally obtaining 4,300 images for model training and testing.

[0039] The target images in the dataset are annotated using LabelImg annotation software to obtain an XML format file, and then a Python script file is used to convert the XML format file into a YOLO format dataset, that is, converting the XML format file into a txt file, so that the converted dataset is suitable for training and testing of more target detection models.

[0040] S2. Use the MBConv structure in the EfficientNet network and combine it with the CA attention mechanism to build the CSP_CAMBConv structure, improve the activation function in the CA attention mechanism, and finally use this structure to reconstruct the backbone network of Yolov5s.

[0041] Specifically, the MBConv structure and CA attention mechanism in the EfficientNet network are integrated to improve the CSP module in the original Yolov5s backbone network. The attention mechanism module in the MBConv structure will be recombined with the CA attention mechanism, which will improve the target detection accuracy while reducing the number of parameters. CAMBneck, which integrates the attention mechanism, will be used as a feature extraction structure to improve the original CSP module and form a new CSP_CAMBneck feature map extraction module.

[0042] In order to enhance the model's attention to small targets in complex backgrounds, the improved CA attention mechanism is integrated into the MBneck structure. The improvements are as follows: the activation function in the attention mechanism is improved, and a smoother and less computationally expensive activation function is proposed based on the Gaussian error function. In order to simplify the computational complexity of the GELU activation function and improve computational efficiency, the error function erf is used to improve the original tan function. The improved activation function FlexGELU formula is as follows:

[0043]

[0044] Among them, erf(x) is the error function, which provides smooth and continuous output based on the integral of the normal distribution. x is the input feature map. In the improved CA attention mechanism, the input in FlexGELU(x) comes from the spatial information and channel information calculated by the CA mechanism. The simplified Gaussian activation function can enhance the nonlinear expression ability of the model and improve the modeling ability of complex features.

[0045] like Figure 2 As shown in the figure, the structure of CSP_CAMBConv is as follows: the input is divided into two branches after being processed by the CBS module. The first branch is directly passed to the Concat module for feature map splicing. The second branch obtains features through the CAMBConv structure and then passes it to the Concat module. The spliced ​​feature map is further extracted through the CBS module to standardize features.

[0046] The improvement to the original Yolov5s backbone network is that the new structure CSP_CAMBConv is applied to the 4th, 6th, and 8th CSP modules in the backbone network respectively; the MBConv structure in the EfficientNet feature extraction network is applied to the ordinary convolution modules in the backbone network, namely the 2nd, 3rd, 5th, 7th, and 9th layers, forming a new Yolov5s backbone network structure.

[0047] S3. A path aggregation structure is used in the neck network to fuse feature information at different levels, and scale inconsistency terms and width difference terms are introduced into the CIoU loss function to construct the SCIoU loss function.

[0048] Since the target may occupy a small size in the image, the deep network fails to make full use of high-level semantic information. Therefore, based on the idea of ​​multi-scale feature fusion, connecting feature layers of different depths can fully fuse rich information at different levels. The 5th and 7th layer feature maps in the backbone network are input to the 1st and 2nd feature fusion modules in the neck network. The 10th layer feature map is output by the SPPF module as the bottom-up input in the neck network. The output of the 1st CBS module and the output of the 2nd CBS module in the neck network are used as the input of the 3rd and 4th feature fusion modules in the top-down.

[0049] like Figure 3 As shown in the figure, the three layers of feature maps extracted by the backbone network are processed. Since the backbone network undergoes multiple feature extractions, the sizes of feature maps at different levels are inconsistent. Therefore, the 10th layer feature map is processed in two ways after being processed by the CBS module. The first way uses upsampling convolution to restore the size of the 10th layer feature map to the same size as the 7th layer feature map, and the 10th layer feature map after upsampling is subjected to the first BiFPN_Concat operation with the 7th layer feature map. The second way performs the fourth BiFPN_Concat operation. The feature map output by the first BiFPN_Concat module is enhanced with feature extraction and fusion by the CSP module. After being processed by the convolution module, the feature map is divided into two ways again for subsequent processing: the first way performs the second BiFPN-Concat operation with the 5th layer feature map through upsampling, and the second way performs the third BiFPN_Concat operation. The feature map output by the second BiFPN_Concat module is processed by CSP and then subjected to the third BiFPN_Concat operation with the result of the second BiFPN_Concat. The two inputs of the fourth BiFPN_Concat operation are the output of the third BiFPN_Concat and the output of the 10th layer feature map after processing by the convolution module. Finally, the CSP outputs of different levels in the Neck layer are sent to the Head layer to obtain object detection results of three different sizes.

[0050] The loss function mainly includes classification loss, confidence loss and bounding box loss. Based on the regression loss function CIoU, the scale inconsistency term and width difference term are introduced to weigh the difference between the predicted box and the real box, and make the model more sensitive to the scale difference of the target during training. The improved SCIoU loss formula is as follows:

[0051]

[0052] Where IoU is the intersection over union ratio between the predicted box and the true box. b is the center point of the predicted bounding box, b gt is the center point of the true bounding box. The Euclidean distance ρ represents the diagonal length of the minimum closed box between the center points of the two bounding boxes. gt 、h gt and w pred 、h pred Represent the width and height of the true bounding box and the predicted bounding box, respectively. c is the diagonal distance of the minimum bounding rectangle, and ε is a small constant to avoid zero division errors. The scale difference of the target is measured by logarithmically transforming the area ratio of the predicted box and the true box. The logarithmic operation makes the model more sensitive when the scale difference is large. The width difference term can reduce the impact of the width difference on the overall loss function by the ratio of the difference between the maximum and minimum values.

[0053] The improved loss function can fully capture the feature information of small targets, improve the recognition and positioning accuracy of small targets in deep learning models, and work together with the CA attention mechanism to solve the problem of insufficient target recognition ability while enhancing the saliency of targets under complex background interference so that the model can focus on the target. In complex scenarios, the model performs well in dealing with overlapping and close small targets, and the recognition accuracy of targets is improved.

[0054] S4. Perform sparse training on the improved target detection model to facilitate subsequent model pruning, and use the global attenuation method to adjust the sparse factor during the sparse process to quickly achieve the sparseness of the sparse factor. Finally, fine-tune the network to obtain the optimal target detection model.

[0055] like Figure 4 As shown in the figure, the model pruning process is as follows: first, the improved network model is initially trained on the data set to obtain the basic model; then the initial model is sparsely trained to highlight important features; then the local structure of the model is pruned to remove unimportant connections; then the pruned network model is fine-tuned and optimized; finally, the optimized pruned model is obtained.

[0056] Among them, in the BN layer sparse training, L1 regularization is applied to the model parameter λ. In addition, in the model sparse training, the sparse factor is dynamically adjusted by introducing a global attenuation method, so that the regularization weight gradually decreases with the training process, and the parameter γ is quickly sparse. The improved sparse training loss function is shown in the following formula:

[0057]

[0058] In the formula, λ is the sparse coefficient corresponding to the pruning sparsification of the model; l is the training loss function before the model introduces regularization; x represents the training input data; f is the model forward prediction function; y is the true label; t represents the current training round number, T represents the total number of training rounds; Γ represents the scaling factor in the BN layer in the model. The basic linear attenuation term is It is the core attenuation part. As the number of training times increases, the sparse coefficient λ will gradually decrease. Dynamic fluctuation term Controlled by β, greater regularization can be applied at the beginning of training, exerting a greater influence on the sparse coefficient λ in the early and middle stages of training, while gradually reducing the influence in the later stages of training to adapt to the overall attenuation trend.

[0059] Finally, after the model completes channel pruning, the detection performance of the model will decrease slightly due to the reduction in parameters and the change in network structure. Therefore, in order to maximize the performance of the detection network, fine-tuning training is required to compensate for the impact of pruning on the accuracy of the detection model. The above steps are iterated continuously to obtain the network model with the best effect after pruning.

[0060] Finally, it should be noted that the above embodiments are only used to illustrate the technical solution of the present invention rather than to limit it. Although the present invention has been described in detail with reference to the preferred embodiments, those skilled in the art should understand that the technical solution of the present invention can be modified or replaced by equivalents without departing from the purpose and scope of the technical solution, which should be included in the scope of the claims of the present invention.

Claims

1. A lightweight object detection method integrating attention mechanism, characterized in that: The method includes: Obtain a public target image dataset and preprocess the dataset, including dataset augmentation, annotation, and partitioning, and finally select the training and test datasets; The MBConv structure in the EfficientNet network is used, combined with the CA attention mechanism to construct the CSP_CAMBConv structure, and the activation function in the CA attention mechanism is improved. The CSP_CAMBConv structure is used to reconstruct the backbone network of Yolov5s. In the neck network of Yolov5s, a path aggregation structure is used to fuse feature information at different levels, and scale inconsistency terms and width difference terms are introduced into the CIoU loss function to construct the SCIoU loss function; The improved target detection model is trained in a sparse manner to facilitate subsequent model pruning. The sparse factor is adjusted using a global attenuation method during the sparse process to quickly achieve the sparseness of the sparse factor. Finally, the network is fine-tuned to obtain the optimal target detection model. Acquire the image and input it into the optimal object detection model for object detection.

2. The method according to claim 1, characterized in that In the CSP_CAMBConv structure, the input is first processed by a CBS module and then passed to the Concat module and the CAMBConv structure respectively. The feature map is spliced ​​by the Concat module, and the features are acquired by the CAMBConv structure and then passed to the Concat module. The spliced ​​feature map is further extracted with standardized features by another CBS module.

3. The method according to claim 2, characterized in that The constructed CSP_CAMBConv is applied to the Yolov5s backbone network to improve Yolov5s; wherein, CSP_CAMBConv is applied to the 4th, 6th and 8th CSP modules of the Yolov5s backbone network; in addition, the MBConv structure in the EfficientNet feature extraction network is applied to the ordinary convolution modules in the backbone network, namely the 2nd, 3rd, 5th, 7th and 9th layers, to form a new Yolov5s backbone network structure.

4. The method according to claim 1, characterized in that: The improved activation function in the CA attention mechanism includes changing the tan function in the original GELU activation function to the erf function to obtain an improved activation function: Where erf(x) represents the error function and x represents the input feature map.

5. The method according to claim 1, characterized in that The neck network of Yolov5s adopts a path aggregation structure to fuse feature information of different levels, including inputting the feature maps output by the 5th and 7th layers in the Yolov5s backbone network into the first and second feature fusion modules in the neck network respectively, and the feature map output by the 10th layer of the backbone network is output to the first CBS module in the neck network after passing through the SPPF module; the outputs of the first CBS module and the second CBS module in the neck network are respectively used as inputs of the third and fourth feature fusion modules in the neck network.

6. The method according to claim 1, characterized in that The scale inconsistency term and the width difference term are introduced into the CIoU loss function, and the SCIoU loss function is constructed, including introducing the scale inconsistency term and the width difference term on the basis of the CIoU loss function to weigh the difference between the predicted box and the true box, and making the model more sensitive to the scale difference of the target during the training process. The improved SCIoU loss function is as follows: In the formula, IoU represents the intersection over union ratio between the predicted box and the true box; b represents the center point of the predicted bounding box, and b gt represents the center point of the true bounding box; ρ represents the Euclidean distance; w gt and h gt Represent the width and height of the real bounding box respectively; w pred and h pred denote the width and height of the predicted bounding box respectively; c denotes the diagonal distance of the minimum enclosing rectangle, and ε denotes a tiny constant to avoid division by zero errors.

7. The method according to claim 1, characterized in that The improved Yolov5s is initially trained to obtain a basic model; then the initial model is sparsely trained to highlight important features; then the local structure of the model is pruned to remove unimportant connections; and the pruned network model is fine-tuned and optimized to obtain an optimized pruned model.

8. The method according to claim 1, characterized in that In sparse training, global attenuation is introduced to dynamically adjust the sparse factor, so that the regularization weight gradually decreases with the training process, and the sparsification of the parameter γ is quickly achieved; the improved sparse training loss function is shown in the following formula: In the formula, λ represents the sparse coefficient corresponding to the pruning sparsification of the model; l represents the training loss function before the regularization is introduced into the model; x′ represents the training input data; f represents the model forward prediction function; y represents the true label; t represents the current training round number; T represents the total number of training rounds; Γ represents the scaling factor of the BN layer in the model.

Images