Lightweight target detection model construction method based on YOLOv9

Abstract

The invention discloses a lightweight target detection model construction method, and belongs to the field of computer vision and model compression. According to the method, YOLOv9s is taken as a benchmark, a RepNCSPELAN4 module is replaced by a MobileOneBlock module for training multiple branches and reasoning a single path through structural re-parameterization, a FlashAttention mechanism is introduced into a feature fusion network, and video memory occupation is optimized through block calculation; and converting into a single-path reasoning structure after cooperative training. Experiments show that mAP50-95 on PASCALVOC2012 and DOTAv1.0 data sets is only reduced by 0.3%-0.4%, the calculation amount is reduced by 30.5%, the parameter amount is reduced by 32%, the problems that when an existing method is applied to a YOLO series, the feature fusion capacity is reduced, and the memory overhead is large are solved, and the method is suitable for real-time target detection tasks with limited resources.

Description

Technical Field

[0001] This invention relates to the field of computer vision technology, and in particular to a method for constructing a lightweight object detection model based on the YOLOv9s architecture, which is suitable for real-time object detection tasks in edge computing devices. Background Technology

[0002] Object detection is a core task in computer vision, widely used in autonomous driving, industrial quality inspection, and intelligent security. The YOLO series of models has become a mainstream solution due to its good balance between speed and accuracy. YOLOv9 further improves detection performance by introducing Programmable Gradient Information (PGI) and Generalized Convergence Network (GELAN). However, the YOLOv9 model has a complex structure, high computational cost, and high memory consumption, making it difficult to deploy on edge devices with limited computing resources.

[0003] Existing lightweight methods mainly include: 1. Network architecture design, such as MobileNet and ShuffleNet, which use depthwise separable convolution or channel shuffling; 2. Structural reparameterization techniques, such as RepVGG, improve efficiency through multi-branch training and single-branch inference; 3. Model compression techniques, such as pruning, quantization, and knowledge distillation.

[0004] However, these methods have the following problems when applied to the YOLO series: 1. The existing method 1 (RepVGG) structure reparameterization technique is only applicable to simple classification networks. When it is directly applied to YOLOv9s, it will destroy the multi-scale feature fusion capability of the RepNCSPELAN4 module, resulting in a decrease of more than 5% in mAP50-95. 2. Existing technology 2 (FlashAttention) only optimizes memory usage, but it is not designed in conjunction with structural reparameterization. When applied to YOLOv9s alone, the computational load is reduced by less than 20%, which cannot meet the needs of edge devices.

[0005] Therefore, how to significantly reduce the computational load and parameter count of the model while maintaining detection accuracy has become a pressing technical problem to be solved in this field. This invention achieves an accuracy loss of <0.5% and a computational load reduction of >30% through a collaborative design that maintains accuracy by training multiple branches and significantly reduces GPU memory usage during the training phase using single-path and block attention. Summary of the Invention

[0006] This invention provides a lightweight object detection model construction method for edge devices, belonging to the field of computer vision and model compression. The method includes the following steps: S1: The backbone network of the benchmark detection model is reconstructed in a lightweight manner using structural reparameterization technology, replacing the standard convolutional modules in the original model with lightweight modules that have multiple branches during training and a single path during inference. S2: Introduce a memory-optimized attention mechanism in the feature fusion network, and reduce the memory usage during attention computation through block computation and kernel fusion strategies; S3: Perform end-to-end co-training on the improved model and optimize the model parameters using the object detection dataset; S4: After training, the model is converted into a single-path inference structure suitable for efficient deployment through structural reparameterization. S5: Evaluate the accuracy, computational cost, number of parameters, and inference speed of the lightweight model on a public object detection dataset to verify its effectiveness.

[0007] Furthermore, step S1 includes: S11: Using the YOLOv9s model as the benchmark detection model, identify all RepNCSPELAN4 modules in its neck network; S12: Replace the standard convolutional components inside the RepNCSPELAN4 module with the MobileOneBlock module, which consists of a depthwise separable convolutional branch, an identity mapping branch, and a 1×1 convolutional branch in parallel during the training phase. S13: By replacing the model, the theoretical computational complexity and number of parameters of the model are significantly reduced while maintaining the model's feature extraction capability.

[0008] Furthermore, step S2 includes: S21: Select key feature fusion nodes in the neck network of the model and replace the standard self-attention mechanism with the FlashAttention mechanism; S22: The FlashAttention mechanism avoids explicitly constructing an N×N attention matrix by dividing the input sequence into blocks and using incremental Softmax computation and kernel fusion technology; S23: Through the aforementioned optimization, the memory complexity of attention computation is reduced from O(N²) to O(N), enabling it to handle higher resolution input features.

[0009] Furthermore, step S3 includes: S31: Jointly train the improved model using the SGD optimizer and the OneCycleLR learning rate scheduling strategy; S32: Apply Mosaic and MixUp data augmentation strategies during training to improve the model's generalization ability; S33: Set the weight of the bounding box regression loss to be higher than that of the classification loss to enhance the model's learning of target localization accuracy.

[0010] Furthermore, step S4 includes: S41: After the model training is completed, perform a structural reparameterization operation on all MobileOneBlock modules; S42: Mathematically equivalently merges multiple parallel branches in each MobileOneBlock module into a single 3×3 standard convolutional layer; S43: Export the fused model to ONNX or TensorRT format for optimal inference performance on edge computing devices.

[0011] Furthermore, step S5 includes: S51: Test the mean accuracy (mAP) of the lightweight model on the PASCAL VOC2012 dataset and the DOTAv1.0 dataset, and compare it with the original benchmark model; S52: Calculate and compare the computational cost (GFLOPs) and parameter count (Params) of the model before and after the improvement. Experimental results show that the lightweight target detection model constructed in this invention reduces the computational load by 30.5% and the number of parameters by 32% with minimal loss in accuracy (mAP@0.5-0.95 decreases by only 0.4%), providing an effective solution for real-time target detection tasks in resource-constrained environments. Attached Figure Description

[0012] Figure 1 This is a schematic diagram of the overall architecture of the lightweight target detection model of the present invention, wherein: 13, 34, 38, 42, 16, 20, and 24 are RepNCSPELAN4 modules (after being replaced by MobileOneBlock), 35, 39, 43, 17, 21, and 25 are FlashAttention modules, and the rest are YOLOv9s original convolution, concatenation, and other components, illustrating the connection relationship between the modules; Figure 2 This is a structural diagram of the MobileOneBlock module, showing the parallel structure of the depthwise separable convolution branch, the identity mapping branch, and the 1×1 convolution branch during the training phase. Figure 3 This is a schematic diagram of the internal structure of the RepNCSPELAN4_MobileOne module, illustrating the feature flow after the standard convolution is replaced by MobileOneBlock; Figure 4This is a flowchart of the FlashAttention module's workflow, where Q, K, and V represent the query matrix, key matrix, and value matrix, respectively, illustrating the steps of block computation and incremental Softmax. Detailed Implementation

[0013] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be described in detail below with reference to the accompanying drawings and specific embodiments. It should be understood that the specific embodiments described herein are merely illustrative of the invention and are not intended to limit it. All other embodiments obtained by those skilled in the art based on the embodiments of this invention without inventive effort are within the scope of protection of this invention.

[0014] Unless otherwise defined, all technical and scientific terms used herein have the same meaning as commonly understood by one of ordinary skill in the art to which this invention pertains. The terminology used herein is for the purpose of describing particular embodiments only and is not intended to limit the scope of the invention.

[0015] like Figure 1 As shown, this embodiment of the invention provides a method for constructing a lightweight target detection model based on YOLOv9s, including the following steps: S1: The neck RepNCSPELAN4 module of the YOLOv9s model is restructured using structural reparameterization technology, and the standard convolution is replaced with the MobileOneBlock module. S2: Introduce the FlashAttention mechanism into the feature fusion network to reduce the GPU memory usage during attention computation through block computation and memory optimization strategies; S3: Perform end-to-end co-training of the model using an improved loss function and data augmentation strategies; S4: Convert the trained multi-branch model into a single-path inference structure through structural reparameterization. S5: Evaluate the accuracy, computational cost, and inference speed of the lightweight model on multiple object detection datasets; S6: Deploy the optimized model to edge computing devices to verify its performance in real-world applications.

[0016] The experimental environment consisted of an NVIDIA RTX 4090 GPU (24G), PyTorch 1.13, and CUDA 11.8. The initial learning rate was 0.01, the number of training epochs was 300, and the batch size was 16. Before training, the input images were uniformly resized to 640×640 and normalized to the range of [0,1]. The random horizontal flipping probability was 0.5.

[0017] The working principle of this technical solution is as follows: First, taking the YOLOv9s model as a benchmark, we analyze the computationally intensive RepNCSPELAN4 module in its network structure. These modules are responsible for multi-scale feature extraction, but they contain a large number of standard convolution operations, resulting in high computational complexity. Structural analysis shows that these modules account for more than 80% of the total computation of the model and are the key targets for optimization.

[0018] During module optimization, the standard convolution in the RepNCSPELAN4 module was replaced with the MobileOneBlock module. MobileOneBlock employs a multi-branch design during training, including a depthwise separable convolution branch, an identity mapping branch, and a 1×1 convolution branch. The depthwise separable convolution decomposes the standard convolution into two steps: depthwise convolution and pointwise convolution, significantly reducing theoretical computational complexity. The identity mapping branch ensures direct gradient backpropagation, mitigating the vanishing gradient problem. The 1×1 convolution branch enhances the model's feature representation capabilities. This multi-branch design provides rich gradient paths during training, thereby enhancing the model's representational power.

[0019] For feature fusion optimization, a FlashAttention mechanism is introduced at the key neck node of the model. FlashAttention divides the input sequence into multiple sub-blocks using a block-based computation strategy, processing them block by block in the GPU's SRAM, avoiding the explicit construction of the complete attention matrix. This method employs incremental Softmax computation, dynamically updating runtime statistics to achieve linear memory complexity. Simultaneously, kernel fusion technology merges multiple computational steps, reducing memory accesses and improving computational efficiency.

[0020] During the model training phase, a collaborative optimization strategy was employed, using the SGD optimizer and OneCycleLR learning rate scheduling, along with Mosaic and MixUp data augmentation. The weights of the loss function were specifically adjusted, setting the bounding box regression loss weight to 7.5 and the classification loss weight to 0.5 to enhance the model's learning of target localization accuracy.

[0021] After training, the multi-branch structure of MobileOneBlock is fused into a single 3×3 convolutional layer through structural reparameterization. Specifically, this involves mathematically fusing the weights of the depthwise separable convolution with the identity mapping branch, and merging the parameters of the 1×1 convolution branch into the main convolution, thereby achieving the optimization goal of multi-branch training and single-path inference.

[0022] One embodiment of the present invention, such as Figure 2 , 3As shown, the lightweight reconstruction of the RepNCSPELAN4 module of the YOLOv9s model using structural reparameterization technology, replacing the standard convolution with the MobileOneBlock module, includes: The four key RepNCSPELAN4 modules in the YOLOv9s model were analyzed, and the number of input and output channels for each module was recorded. Replace the standard 3×3 convolution in the module with the MobileOneBlock module while maintaining consistent input and output dimensions; During the training phase, the MobileOneBlock module performs the computation of three branches in parallel: depthwise separable convolution, identity mapping, and 1×1 convolution. During the inference phase, the multi-branch structure is transformed into a single 3×3 standard convolutional layer through parameter fusion.

[0023] The working principle of the above technical solution is as follows: First, the network structure of the YOLOv9s model is analyzed, and four key RepNCSPELAN4 modules are identified. These modules are distributed in feature layers at different scales and are responsible for extracting and fusing multi-resolution features. Through computational analysis, it is determined that these modules contain a total of 12 standard 3×3 convolutional layers, resulting in high computational density.

[0024] During module replacement, the consistency of input and output channel numbers is maintained. For convolution operations with an input feature map size of H×W×C_in and an output channel number of C_out, the replaced MobileOneBlock module maintains the same dimensions. Specifically, the depthwise separable convolution branch first performs a 3×3 depthwise convolution, keeping the output channel number at C_in, and then performs a 1×1 pointwise convolution to expand the channel number to C_out. The identity mapping branch directly passes features when the input and output channel numbers are the same; when the channel numbers are different, the dimensions are adjusted through a 1×1 convolution.

[0025] The mathematical representation of the multi-branch structure during the training phase is: in This represents depthwise separable convolution. This represents a 1×1 convolution. This design ensures that the model has strong feature representation capabilities during training. Depthwise separable convolution reduces computational complexity, identity mapping improves gradient flow, and 1×1 convolution enhances feature transformation capabilities.

[0026] The structure reparameterization during the inference phase includes the following steps: First, the depthwise convolution weights of the depthwise separable convolutions are combined with the pointwise convolution weights to obtain an equivalent 3×3 convolution kernel. Second, the identity mapping branches are converted into their corresponding convolution kernels. Finally, the convolution kernels and bias terms of all branches are superimposed to obtain the final parameters of a single convolutional layer. The specific fusion formula is as follows: in These are depthwise convolution weights. is the pointwise convolution weight, and I is the convolution kernel corresponding to the identity matrix.

[0027] The effects of the above technical solution are as follows: by replacing the MobileOneBlock module, the model computational cost was reduced from the baseline of 39.7 GFLOPs to 27.5 GFLOPs, a reduction of 30.7%. The number of parameters was reduced from 9.7M to 7.0M, a reduction of 28.9%. Tests on the PASCAL VOC dataset show that when using this optimization alone, mAP@0.5-0.95 decreased from 56.4% to 51.3%, leaving ample room for improvement in subsequent attention mechanism optimizations.

[0028] One embodiment of the present invention, such as Figure 4 As shown, the introduction of the FlashAttention mechanism into the feature fusion network reduces the GPU memory usage during attention computation through block computation and memory optimization strategies, including: The self-attention computation module is located at the top node of the feature pyramid network; Replace the standard self-attention mechanism with the FlashAttention module, configuring the same query, key, and value matrix dimensions; Implement a block-based computation strategy to divide the input sequence into sub-blocks of size B; An incremental Softmax calculation method is used to avoid storing the complete attention matrix; Optimize GPU computing processes through core fusion technology.

[0029] The working principle of the above technical solution is as follows: First, three key feature fusion nodes are identified in the neck network of YOLOv9s. These nodes are responsible for processing high-level semantic features and have a large receptive field. Standard self-attention mechanisms incur significant memory overhead at these nodes, especially when processing high-resolution feature maps.

[0030] The implementation of FlashAttention consists of two core parts: block-based computation and memory optimization. Block-based computation divides the input sequence into multiple sub-blocks of size B. The size of B is optimized based on the GPU's SRAM capacity, typically between 64 and 256. During the computation process, a query block and its corresponding key-value block are loaded into SRAM at each step for local attention computation.

[0031] The incremental Softmax calculation employs a dynamic update strategy, maintaining the maximum value and summation statistics during operation. The specific algorithm flow is as follows: The initial maximum value is negative infinity, and the summation value is 0. For each data block: Calculate the query-key dot product of the current block; Update the maximum running value; Calculate the exponent value and normalize it; Update the weighted sum output; The final output is the complete attention result.

[0032] Memory access optimization is achieved through kernel fusion, which combines steps such as query partitioning, key-value partitioning, dot product calculation, softmax, and weighted summation into a single GPU kernel function. This optimization reduces multiple accesses to global memory and improves computational efficiency.

[0033] In its implementation, the FlashAttention module is integrated into the neck network of YOLOv9s, replacing the original self-attention component. The input feature map is first flattened into a sequence, normalized by a LayerNorm layer, and then linearly transformed to obtain the query, key, and value matrices. After FlashAttention is calculated, the output features are restored to their spatial dimensions and residually connected to the original feature map.

[0034] Table 1 Ablation Experiments using the VOC2012 Dataset YOLOv9-s MobileOneBlock FlashAttention mAP@0.5 mAP@0.5-0.95 Number of parameters (M) Computational complexity (G) √ 71.5 56.4 9.7 39.7 √ √ 64.7 51.3 7.0 27.5 √ √ 72.8 58.9 7.9 32.1 √ √ √ 70.6 55.0 6.3 27.6 Table 2 Ablation Experiments on the DOTAV1.0 Dataset YOLOv9-s MobileOneBlock FlashAttention mAP@0.5 mAP@0.5-0.95 Number of parameters (M) Computational complexity (G) √ 65.0 45 9.7 39.7 √ √ 63.7 42.2 7.0 27.5 √ √ 66.1 46.7 7.9 32.1 √ √ √ 65.1 44.7 6.5 27.7 To ensure the fairness of the experimental results, all experiments in this invention used the same training metrics. The effects of the above technical solution are as follows: when training the Pascal VOC2012 dataset, the memory usage was reduced from 16GB for standard attention to 6.4GB, a reduction of approximately 60%. Inference speed was improved by 2-3 times, and processing time on an RTX 4090 was reduced from 8.2ms to 3.1ms. When FlashAttention was used alone for optimization, mAP@0.5-0.95 on the VOC dataset increased from 56.4% to 58.9%, showing a good feature enhancement effect. When used in conjunction with MobileOneBlock optimization, the final model achieved a significant reduction in computational cost of 30.5% and parameter count of 32% with minimal loss in accuracy (mAP@0.5-0.95 decreased by only 0.4%). Similarly, on the DOTAv1.0 dataset, a significant reduction in parameter count and computational cost was achieved (parameter count reduced by 3.4M and computational cost reduced by 12.1G), with only a 0.2% loss in accuracy. This demonstrates the effectiveness of the improved solution.

[0035] Table 3. Comparison of different models on the VOC2012 dataset Network Model mAP@0.5 mAP@0.5-0.95 Number of parameters (M) Computational complexity (G) Inference speed (ms) YOLOv8-s 66.7 50.6 11.1 28.7 1.9 YOLOv11-s 63.9 47.9 9.1 24.6 1.2 YOLOv12-s 68.0 52.9 9.1 19.4 1.9 Model in the text 70.6 55.0 6.3 27.6 1.0 Table 4 Comparison Experiments on the DOTAV1.0 Dataset Network Model mAP@0.5 mAP@0.5-0.95 Number of parameters (M) Computational complexity (G) Inference speed (ms) YOLOv8-s 60.9 40.2 11.2 28.1 1.8 YOLOv11-s 62.3 42.9 9.1 24.6 1.4 YOLOv12-s 64 43.7 9.1 19.4 1.6 Model in the text 65.1 44.7 6.5 27.7 1.2

Claims

1. A method for constructing a lightweight target detection model, characterized in that, The method includes: The structural reparameterization step is used to lightweight reconstruct the feature extraction module in the benchmark detection model, replacing the standard convolution with a reparameterizable lightweight module. The attention optimization step involves memory optimization of the attention computation module in the feature fusion network to reduce the GPU memory usage during the attention computation process. The collaborative training step involves end-to-end training and optimization of the reconstructed model; The model conversion step converts the trained multi-branch model into a single-path inference structure.

2. The lightweight target detection model construction method according to claim 1, characterized in that, The structure reparameterization step includes: identifying the RepNCSPELAN4 module in the baseline YOLOv9s model; replacing the standard convolutional components in the RepNCSPELAN4 module with the MobileOneBlock module, which includes a depthwise separable convolutional branch, an identity mapping branch, and a 1×1 convolutional branch during the training phase.

3. The lightweight target detection model construction method according to claim 1 or 2, characterized in that, The attention optimization steps include: locating the self-attention computation position at the key nodes of the feature pyramid network; replacing the standard self-attention mechanism with the FlashAttention mechanism; and dividing the input sequence into multiple sub-blocks for incremental computation using a block computation strategy.

4. The lightweight target detection model construction method according to claim 3, characterized in that, The block-based computation strategy includes: dividing the query matrix, key matrix, and value matrix into sub-blocks of size B in the sequence dimension; performing attention computation block by block in the GPU's SRAM; and using an incremental Softmax computation method to avoid storing the complete attention matrix.

5. The lightweight target detection model construction method according to claim 2, characterized in that, During the inference phase, the MobileOneBlock module uses structural reparameterization technology to merge depthwise separable convolutional branches, identity mapping branches, and 1×1 convolutional branches into a single 3×3 convolutional layer.

6. The lightweight target detection model construction method according to claim 5, characterized in that, The structure reparameterization technique includes: synthesizing depthwise convolution weights and pointwise convolution weights; converting identity mapping branches into corresponding convolution kernels; and superimposing the convolution kernels and bias terms of all branches to obtain the final single convolutional layer parameters.

7. The lightweight target detection model construction method according to claim 1, characterized in that, The collaborative training steps include: using the SGD optimizer in conjunction with the OneCycleLR learning rate scheduling strategy; and employing Mosaic and MixUp data augmentation techniques.

8. The lightweight target detection model construction method according to claim 7, characterized in that, The method also includes a model evaluation step: testing the model's average accuracy on the PASCAL VOC and DOTA datasets; calculating the computational cost and parameter variations of the model; and testing the model's inference speed on an edge computing platform.

9. A lightweight target detection system, characterized in that: The model reconstruction module is used to perform lightweight structural modifications on the benchmark detection model. The attention optimization module is used to optimize the memory usage of attention computation in the feature fusion network. The training optimization module is used for co-training the improved model; Deploy the conversion module to convert the trained model into an inference format suitable for edge devices.

10. The lightweight target detection system according to claim 9, characterized in that, The model reconstruction module specifically includes: a module identification unit for identifying computationally intensive modules in the baseline model; a structure replacement unit for replacing standard convolution with MobileOneBlock modules; and a parameter fusion unit for fusing multi-branch structures into a single-path inference structure.

Images