A lightweight object detection method integrating attention mechanism

Abstract

This invention relates to a lightweight object detection method incorporating an attention mechanism, belonging to the field of image recognition technology. The method first creates and preprocesses a target image dataset; then, it constructs a CSP_CAMBConv structure by combining the MBConv structure and the CA attention mechanism, and uses this structure to reconstruct the Yolov5s backbone network. A bidirectional feature pyramid fusion module is employed in the neck network to enhance feature information transmission and sharing; furthermore, scale inconsistency and width difference terms are introduced into the loss function to balance the difference between the predicted and ground truth bounding boxes, making the model more sensitive to target scale differences during training. Finally, sparse training is used to prune and compress the network model size, thereby obtaining the optimal lightweight model. This invention is suitable for edge terminals with limited computing resources, enabling accurate localization of targets of different scales in complex backgrounds, effectively reducing false detections and false negatives.

Description

Technical Field

[0001] This invention belongs to the field of image recognition terminal technology and relates to a lightweight target detection method that integrates an attention mechanism. Background Art

[0002] In recent years, deep neural networks based on convolutional operations have achieved superior performance in many fields of machine vision and natural language processing due to their excellent generalization ability and efficient feature extraction capabilities. Machine vision object detection algorithms can quickly locate and classify specific targets, while convolutional neural network-based object detection technology is widely used in military, industrial quality inspection, and other fields. Although deep neural networks have wide applications, there is still considerable room for improvement in engineering deployment. Deep neural networks have high time and space complexity, requiring significant computational and storage resources, placing high demands on hardware computing and storage resources, making them difficult to deploy on resource-constrained terminals. Therefore, to ensure that object detection algorithms can run under resource constraints, it is necessary to lightweight the detection network.

[0003] Furthermore, the Internet of Things (IoT), as a significant derivative of the Internet, provides interconnectivity and intelligent transmission capabilities. With the proliferation of mobile terminals, target detection models are gradually becoming miniaturized and shifting from traditional server deployments to edge terminal platforms. However, improving target detection speed and accuracy remains a critical challenge given the limited computing resources at the edge. Therefore, reducing the number of network model parameters and computational load while maintaining target detection accuracy, and deploying target recognition algorithms stably on resource-constrained hardware platforms, has practical application value and broad engineering application prospects. Summary of the Invention

[0004] In view of this, the purpose of this invention is to provide a lightweight object detection method that integrates an attention mechanism to solve the problem of deploying a high-accuracy object detection model on a terminal with limited computing resources.

[0005] To achieve the above objectives, the present invention provides the following technical solution:

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

[0007] S1. Obtain the publicly available target image dataset and preprocess the dataset, including dataset augmentation, annotation, and partitioning, and finally select the training and testing datasets.

[0008] S2. The BMBConv structure in the EfficientNet network is adopted, and the CSP_CAMBConv structure is constructed by combining the CA attention mechanism. The activation function in the CA attention mechanism is improved, and the Yolov5s backbone network is reconstructed using the CSP_CAMBConv structure.

[0009] S3. In the neck network of YOLOv5s, a path aggregation structure is adopted to fuse feature information from different levels, and scale inconsistency term and width difference term are introduced into CIoU loss function to construct SCIoU loss function;

[0010] S4. The improved target detection model is trained in a sparsity manner to facilitate subsequent model pruning. During the sparsity process, the sparsity factor is adjusted by global decay to quickly achieve sparsity factor reduction. Finally, the network is fine-tuned to obtain the optimal lightweight target detection model.

[0011] S5. Acquire the image and input it into the optimal object detection model for object 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 Concat module performs feature map concatenation, and the CAMBConv structure obtains features before passing them to the Concat module. The concatenated feature map is then further processed by another CBS module to extract standardized features.

[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 CSP modules of layers 4, 6, and 8 of the Yolov5s backbone network. In addition, the MBConv structure in the EfficientNet feature extraction network is applied to the ordinary convolutional modules in the backbone network, namely layers 2, 3, 5, 7, and 9, to form a new Yolov5s backbone network structure.

[0014] Furthermore, in step S2, the activation function in the CA attention mechanism is improved by changing the tan function in the original GELU activation function to the erf function, resulting in the improved activation function:

[0015]

[0016] In the formula, 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 from different levels. Specifically, the feature maps output from layers 5 and 7 of the YOLOv5s backbone network are input to the first and second feature fusion modules in the neck network, respectively. The feature map output from layer 10 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 and second CBS modules in the neck network are used as the inputs to the third and fourth feature fusion modules in the neck network, respectively.

[0018] Furthermore, in step S3, a scale inconsistency term and a width difference term are introduced into the CIoU loss function to construct the SCIoU loss function. Specifically, a scale inconsistency term and a width difference term are introduced into the CIoU loss function to balance the difference between the predicted bounding box and the ground truth bounding box, and to make the model more sensitive to the scale difference of the target during training. The improved SCIoU loss function is shown below:

[0019]

[0020] In the formula, IoU represents the intersection-union ratio between the predicted bounding box and the ground truth bounding box; b represents the center point of the predicted bounding box. gt The center point of the true bounding box is represented by ρ; the Euclidean distance is represented by w. gt and h gt These represent the width and height of the actual bounding box, respectively; w pred and h pred Let 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 division-by-zero errors. The improved SCIoU loss function measures the scale difference of the target by performing a logarithmic transformation on the ratio of the areas of the predicted and ground truth boxes. The logarithmic operation makes the model more sensitive to large scale differences. The width difference term, through the ratio of the difference between the maximum and minimum values, can reduce the impact of width differences on the overall loss function.

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

[0022] Meanwhile, in the sparsity training, a global decay is introduced to dynamically adjust the sparsity factor, so that the regularization weight gradually decreases during the training process, quickly achieving sparsity of parameter γ; the improved sparse training loss function is shown in the following equation:

[0023]

[0024] In the formula, λ represents the sparsity coefficient corresponding to model pruning and sparsification; l represents the training loss function before regularization; x′ represents the training input data; f represents the model's forward prediction function; y represents the true label; t represents the current training epoch; T represents the total number of training epochs; and Γ represents the scaling factor of the BN layer in the model. The basic linear decay term... This is the core decay component; as the number of training iterations increases, the sparsity coefficient λ will gradually decrease. Dynamic fluctuation term. Controlled by β, greater regularization can be applied in the early stages of training, exerting a greater influence on the sparse coefficients λ in the early and middle stages of training, while gradually reducing the influence in the later stages of training to adapt to the overall decay trend.

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

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

[0027] (2) To meet the detection requirements of targets at different scales, this 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 balance the difference between the predicted box and the real box, and to make the model more sensitive to the scale difference of the target during the training process.

[0028] (3) This 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 invention will be set forth in part in the description which follows, and in part will be apparent to those skilled in the art from the following examination, or may be learned from practice of the invention. The objectives and other advantages of the invention can be realized and obtained through the following description. Attached Figure Description

[0030] To make the objectives, technical solutions, and advantages of the present invention clearer, the preferred embodiments of the present invention will be described in detail below with reference to the accompanying drawings, wherein:

[0031] Figure 1 This is a flowchart illustrating the method described in this invention;

[0032] Figure 2This is a schematic diagram of the CSP_CAMBConv structure proposed in an embodiment of the present invention;

[0033] Figure 3 This is a diagram of a lightweight network structure proposed in an embodiment of the present invention;

[0034] Figure 4 Flowchart for pruning the model. Detailed Implementation

[0035] The following specific examples illustrate the implementation of the present invention. Those skilled in the art can easily understand other advantages and effects of the present invention from the content disclosed in this specification. The present invention can also be implemented or applied through other different specific embodiments, and various details in this specification can be modified or changed 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 are only schematic representations of the basic concept of the present invention. Unless otherwise specified, the following embodiments and features can be combined with each other.

[0036] like Figure 1 As shown, this invention provides a lightweight object detection method that integrates an attention mechanism. The specific details of this method are as follows:

[0037] S1. Obtain the publicly available target image dataset and preprocess the dataset, including dataset augmentation, annotation, and partitioning, and finally select the training and testing datasets.

[0038] In this embodiment, the publicly available dataset VisDrone2019 is used as the test, training, and validation dataset. Various data augmentation preprocessing techniques are applied to the dataset to generate more diverse samples, including augmentation using rotation, mirroring, noise reduction, and brightness variations. Invalid images that do not contain the target due to preprocessing are then removed, ultimately yielding 4300 images for model training and testing.

[0039] The target images in the dataset are labeled using LabelImg annotation software to obtain XML format files. Then, a Python script is used to convert the XML format files into YOLO format datasets, that is, to convert the XML format files into TXT files, so that the converted datasets can be used for training and testing of more object detection models.

[0040] S2. The MBConv structure in the EfficientNet network is adopted, and the CSP_CAMBConv structure is constructed by combining it with the CA attention mechanism. The activation function in the CA attention mechanism is improved, and finally the Yolov5s backbone network is reconstructed using this structure.

[0041] Specifically, the MBConv structure and CA attention mechanism from the EfficientNet network are integrated to improve the CSP module in the original Yolov5s backbone network. The attention mechanism module within the MBConv structure will be recombined with the CA attention mechanism, improving object detection accuracy while reducing the number of parameters. The CAMBneck, incorporating the attention mechanism, will serve as a feature extraction structure to improve the original CSP module, forming a new CSP_CAMBneck feature map extraction module.

[0042] To enhance the model's attention to small targets in complex backgrounds, an improved CA attention mechanism is integrated into the MBneck structure. Specifically, the activation function in the attention mechanism is improved by proposing a smoother and more computationally efficient activation function based on the Gaussian error function. Furthermore, to simplify the computational complexity of the GELU activation function and improve computational efficiency, the original tan function is modified using the error function erf. The improved activation function, FlexGELU, is shown in the following formula:

[0043]

[0044] Here, erf(x) is the error function, based on the integral of the normal distribution, providing a smooth and continuous output. x is the input feature map. In the improved CA attention mechanism, the input to FlexGELU(x) comes from the spatial and channel information calculated by the CA mechanism. The simplified Gaussian activation function can enhance the nonlinear expressive power of the model and improve its ability to model complex features.

[0045] like Figure 2 As shown, the specific structure of CSP_CAMBConv is as follows: After the input is processed by the CBS module, it is divided into two branches. The first branch is directly passed to the Concat module for feature map concatenation. The second branch obtains features through the CAMBConv structure and then passes them to the Concat module. The concatenated feature map is further processed by the CBS module to extract standardized features.

[0046] The improvement to the original Yolov5s backbone network is achieved by applying the new structure CSP_CAMBConv to the CSP modules in layers 4, 6, and 8 of the backbone network, respectively; and by applying the MBConv structure from the EfficientNet feature extraction network to the ordinary convolutional modules in the backbone network, namely layers 2, 3, 5, 7, and 9, to form a new Yolov5s backbone network structure.

[0047] S3. In the neck network, a path aggregation structure is used to fuse feature information from different levels, and scale inconsistency term and width difference term 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, deep networks fail to fully utilize high-level semantic information. Therefore, based on the idea of ​​multi-scale feature fusion, connecting feature layers of different depths can fully integrate rich information from different levels. The feature maps of layers 5 and 7 in the backbone network are input to the first and second feature fusion modules in the neck network. The feature map of layer 10 is output by the SPPF module as the bottom-up input in the neck network. The outputs of the first and second CBS modules in the neck network are input to the third and fourth feature fusion modules in the top-down approach.

[0049] like Figure 3 As shown, the three-layer feature maps extracted by the backbone network are processed. Because the backbone network undergoes multiple feature extractions, the sizes of the feature maps at different levels are inconsistent. Therefore, the 10th-layer feature map is processed in two paths after being processed by the CBS module. The first path uses upsampling convolution to restore the size of the 10th-layer feature map to the same size as the 7th-layer feature map. The upsampled 10th-layer feature map is then combined with the 7th-layer feature map for the first BiFPN_Concat operation. The second path performs the fourth BiFPN_Concat operation. The feature map output from the first BiFPN_Concat module is enhanced for feature extraction and fusion using the CSP module. After processing by the convolution module, this feature map is again divided into two paths for subsequent processing: the first path performs a second BiFPN_Concat operation with the 5th-layer feature map through upsampling, and the second path performs a third BiFPN_Concat operation. The feature map output from the second BiFPN_Concat module, after being processed by CSP, is then combined with the result of the second BiFPN_Concat operation for the third BiFPN_Concat operation. The four BiFPN_Concat operations take two inputs: the output of the third BiFPN_Concat operation and the output of the 10th layer feature map after convolutional processing. Finally, the CSP outputs from different levels in the Neck layer are fed into 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, a scale inconsistency term and a width difference term are introduced to balance the difference between the predicted and ground truth boxes, making the model more sensitive to the scale differences of the target during training. The improved SCIoU loss formula is as follows:

[0051]

[0052] In the formula, IoU measures the intersection-union ratio between the predicted bounding box and the ground truth bounding box. b is the center point of the predicted bounding box. gt Let be the center point of the true bounding box. The Euclidean distance ρ represents the diagonal length of the smallest closed box between the center points of two bounding boxes. And w... gt h gt and w pred h pred represents the width and height of the ground truth 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 division by zero errors. The formula measures the scale difference of the target by applying a logarithmic transformation to the ratio of the predicted and ground truth bounding boxes' areas. The logarithmic operation makes the model more sensitive to large scale differences. The width difference term, through the ratio of the difference between the maximum and minimum values, can reduce the impact of width differences on the overall loss function.

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

[0054] S4. The improved target detection model is trained with sparsity to facilitate subsequent model pruning. During the sparsity process, the sparsity factor is adjusted using a global decay method to quickly achieve sparsity of the sparsity factor. Finally, the network is fine-tuned to obtain the optimal target detection model.

[0055] like Figure 4 As shown, the model pruning process is as follows: First, the improved network model is initially trained on the dataset to obtain the basic model; then, the initial model is sparsely trained to highlight important features; next, 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] In the sparsity training of the BN layer, L1 regularization is applied to the model parameter λ. Furthermore, during model sparsity training, a global decay mechanism is introduced to dynamically adjust the sparsity factor, causing the regularization weights to gradually decrease during training, thus rapidly achieving sparsity of parameter γ. The improved sparse training loss function is shown in the following formula:

[0057]

[0058] In the formula, λ is the sparsity coefficient corresponding to model pruning and sparsification; l is the training loss function before regularization; x represents the training input data; f is the model's forward prediction function; y is the true label; t represents the current training epoch, T represents the total number of training epochs; and Γ represents the scaling factor in the BN layer of the model. The basic linear decay term... This is the core decay component; as the number of training iterations increases, the sparsity coefficient λ will gradually decrease. Dynamic fluctuation term. Controlled by β, greater regularization can be applied in the early stages of training, exerting a greater influence on the sparse coefficients λ in the early and middle stages of training, while gradually reducing the influence in the later stages of training to adapt to the overall decay trend.

[0059] Finally, after channel pruning, the model's detection performance will slightly decrease due to the reduction in parameters and the change in network structure. Therefore, to maximize the performance of the detection network, fine-tuning is needed to compensate for the impact of pruning on the accuracy of the detection model. By iterating through these steps, the optimal network model after pruning can be obtained.

[0060] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and are not intended to limit it. Although the present invention has been described in detail with reference to preferred embodiments, those skilled in the art should understand that modifications or equivalent substitutions can be made to the technical solutions of the present invention without departing from the spirit and scope of the present invention, and all such modifications or substitutions should be covered within the scope of the claims of the present invention.

Claims

1. A lightweight object detection method incorporating an attention mechanism, characterized in that, The method includes: Obtain publicly available target image datasets and preprocess the datasets, including dataset augmentation, annotation, and partitioning, and finally select training and testing datasets. We adopt the MBConv structure from the EfficientNet network, combine it with the CA attention mechanism to construct 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. In the neck network of YOLOv5s, a path aggregation structure is used to fuse feature information from different levels, and scale inconsistency term and width difference term are introduced into CIoU loss function to construct SCIoU loss function; The improved target detection model is trained with sparsity to facilitate subsequent model pruning. During the sparsification process, the sparsity factor is adjusted using a global decay method to quickly achieve sparsity of the sparsity factor. Finally, the network is fine-tuned to obtain the optimal target detection model. Acquire images and input them into the optimal object detection model for object detection; 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 Concat module performs feature map concatenation, and the CAMBConv structure obtains features and then passes them back to the Concat module. The concatenated feature map is further processed by another CBS module to extract standardized features. The constructed CSP_CAMBConv structure is applied to the YOLOv5s backbone network to improve YOLOv5s; specifically, the CSP_CAMBConv structure is applied to the CSP modules of layers 4, 6, and 8 of the YOLOv5s backbone network; in addition, the MBConv structure from the EfficientNet feature extraction network is applied to ordinary convolutional modules in the backbone network, namely layers 2, 3, 5, 7, and 9, to form a new YOLOv5s backbone network structure. The improved CA attention mechanism includes replacing the tan function in the original GELU activation function with the erf function, resulting in the improved activation function: In the formula, Represents the error function. Represents the input feature map; The method of fusing feature information from different levels in the neck network of YOLOv5s using a path aggregation structure includes: inputting the feature maps output from layers 5 and 7 of the YOLOv5s backbone network into the first and second feature fusion modules BiFPN_Concat in the neck network, respectively; and outputting the feature map output from layer 10 of the backbone network into the first CBS module in the neck network after passing through the SPPF module. The outputs of the first and second CBS modules in the neck network are used as inputs to the third and fourth feature fusion modules BiFPN_Concat in the neck network, respectively. The method of introducing scale inconsistency and width difference terms into the CIoU loss function to construct the SCIoU loss function involves adding scale inconsistency and width difference terms to the CIoU loss function to balance the difference between the predicted bounding box and the ground truth bounding box, and to make the model more sensitive to the scale difference of the target during training. The improved SCIoU loss function is shown below: In the formula, IoU 'b' represents the intersection-union ratio (IoU) between the predicted bounding box and the ground truth bounding box; 'b' represents the center point of the predicted bounding box. This represents the center point of the actual bounding box; Indicates Euclidean distance; and These represent the width and height of the actual bounding box, respectively; and These represent the width and height of the predicted bounding box, respectively. This represents the diagonal distance of the smallest bounding rectangle. This represents a tiny constant that avoids division by zero errors.

2. The method according to claim 1, characterized in that, The improved Yolov5s was initially trained to obtain the base model; then the initial model was sparsified to highlight important features; next, the local structure of the model was pruned to remove unimportant connections; and finally, the pruned network model was fine-tuned to obtain the optimized pruned model.

3. The method according to claim 1, characterized in that, In sparse training, a global decay is introduced to dynamically adjust the sparsity factor, so that the regularization weights gradually decrease during training, quickly achieving parameter optimization. Sparsification; the improved sparse training loss function is shown in the following equation: In the formula, This represents the sparsity coefficient corresponding to the sparsification of the model during pruning; This represents the training loss function before regularization is introduced into the model; This represents the training input data; This represents the model's forward prediction function; Indicates the true label; Indicates the current training round number; Indicates the total number of training rounds; In the representation model The scaling factor of the layer.

Images