A lightweight traffic target detection method based on knowledge distillation

Abstract

The application belongs to the technical field of computer vision, and discloses a light-weight traffic target detection method based on knowledge distillation, wherein a new light-weight detection model paradigm developed by using a knowledge distillation strategy is used to optimize the detection ability of an existing general model on targets in a complex traffic scene; based on the existing light-weight model, a student model and a teacher model comprising a foreground prediction module and a target detection module are established; and a knowledge distillation strategy is designed to improve the student model to be targetedly fused with traffic target features and the expression of traffic target categories and position information at the output end of the model. The application realizes the effective improvement of the detection ability of the existing model on traffic targets by using slight model parameters, and has potential application value in improving the safety of intelligent driving or auxiliary driving and the deployment of a vehicle-mounted platform.

Description

Technical Field

[0001] This invention belongs to the field of computer vision technology and relates to a lightweight traffic target detection method based on knowledge distillation. Background Technology

[0002] The performance of traffic target detection models has a significant impact on the safety of autonomous driving. Developing lightweight models with strong detection capabilities is an effective way to flexibly deploy models on computationally limited in-vehicle platforms. Existing general-purpose lightweight models reduce the number of model parameters by decreasing the model's width and depth, while utilizing multi-scale strategies to improve detection accuracy. However, traffic scenarios are complex and varied, and applying such general-purpose lightweight models to traffic target detection has at least the following shortcomings: Firstly, the lack of targeted guidance during the traffic target feature fusion stage may introduce background noise interference, thus hindering the improvement of model detection performance. Secondly, the limited number of model parameters results in insufficient representation of traffic target features, making it unable to effectively detect targets in complex traffic scenarios. Knowledge distillation, as an effective method for developing lightweight models, can improve model detection performance. The question is how to integrate knowledge distillation into the feature fusion stage of the model for complex traffic scenarios, thereby developing a lightweight detection model paradigm to conveniently and effectively improve the detection performance of existing general-purpose models. Summary of the Invention

[0003] (I) Purpose of the Invention

[0004] The purpose of this invention is to propose a lightweight traffic target detection method based on knowledge distillation. By developing a concise detection model paradigm, it can effectively improve the detection capability of existing lightweight models in complex traffic scenarios, and has potential application value in improving the safety of intelligent driving or assisted driving.

[0005] (II) Technical Solution

[0006] To address the aforementioned technical problems, this invention provides a lightweight traffic target detection method based on knowledge distillation, comprising the following steps:

[0007] S1. Generate a foreground target region mask based on the traffic target detection training set;

[0008] S2. Based on the existing lightweight model using the PANet network structure, a lightweight student model is established. This model mainly includes: a foreground prediction module and an object detection module.

[0009] Among them, the foreground prediction module shares the backbone network and top-down path of the existing lightweight detection model to predict the area where the target is located in the traffic scene;

[0010] The target detection module shares the bottom-up path and detection head of the existing lightweight detection model to characterize the category and location information of traffic targets.

[0011] S3. A lightweight student model built on S2 is constructed, and a teacher model with strong feature expression capabilities is established. This model mainly includes: a foreground prediction module and an object detection module.

[0012] Among them, the foreground prediction module is used to decode rich features related to the traffic target area;

[0013] The target detection module is used to effectively decode the category and location information of traffic targets.

[0014] S4. Establish an effective distillation strategy, which mainly includes: foreground area feature guidance and traffic target location and category information guidance.

[0015] Among them, the rich traffic target area features decoded by the foreground prediction module of the teacher model are used to guide the foreground prediction module of the student model to express the features of the foreground area in the traffic scene.

[0016] The traffic target category and location information decoded by the target detection module of the teacher model are used to guide the target detection module of the student model in expressing the target category and location in the traffic scene.

[0017] S5. Use the traffic target detection training set and the generated foreground target region mask from S1 to train the teacher model, thereby obtaining the network weights;

[0018] S6. The obtained teacher model network weights, the traffic target detection training set in S1, and the generated foreground target region mask are used to jointly supervise the training of the student model, thereby improving the student model's ability to represent traffic targets.

[0019] S7. Save the student model parameters and use the data in the test set to test the detection performance of the lightweight student model.

[0020] (III) Beneficial Effects

[0021] The lightweight traffic target detection method based on knowledge distillation provided by the above technical solution utilizes a new lightweight detection model paradigm developed using knowledge distillation strategies to optimize the detection capability of existing general models in complex traffic scenarios. Based on existing lightweight models, student and teacher models are established, each containing a foreground prediction module and a target detection module. A knowledge distillation strategy is designed to improve the student model's ability to selectively integrate traffic target features and enhance the model's output representation of traffic target category and location information. This invention effectively improves the traffic target detection capability of existing models using minimal model parameters, and has potential application value in enhancing the safety of intelligent driving or assisted driving systems and in vehicle platform deployment. Attached Figure Description

[0022] Figure 1 This refers to the foreground target region mask generated based on the traffic target detection dataset constructed in this embodiment of the invention;

[0023] Figure 2 This is a network structure diagram based on knowledge distillation established in an embodiment of the present invention;

[0024] Figure 3 The diagram shows the detection results of an embodiment of the present invention; where (a) is the input traffic image, (b) is the detection result of the method proposed in this invention, (c) is the Gelan-S detection result, and (d) is the labeled ground truth label. Detailed Implementation

[0025] To make the objectives, contents, and advantages of the present invention clearer, the specific embodiments of the present invention will be described in further detail below with reference to the accompanying drawings and examples.

[0026] This embodiment of the lightweight traffic target detection method based on knowledge distillation includes the following steps:

[0027] S1. To conduct object detection in complex traffic scenarios, the publicly available BDD100K dataset was used for model training, validation, and testing, specifically as follows:

[0028] a) Randomly select 5K images and corresponding labels from the BDD100K training set for the validation set, and use the remaining 65K images and corresponding labels for the new training set. Use the 10K images from the BDD100K validation set as the test set for this case.

[0029] b) To train the foreground prediction module, foreground object region masks are generated based on the object detection labels corresponding to 65K training images. Specifically, the region values ​​within the bounding boxes in the object detection labels are set to 1, and other region values ​​are set to 0. An example of the object detection labels and corresponding region masks is shown below. Figure 1 As shown.

[0030] S2. Based on existing lightweight models using the PANet network structure, such as YOLOV5-S, a lightweight student model STOD-Net is established. This model mainly includes a foreground prediction module (FPM) and an object detection module (ODM).

[0031] The foreground prediction module shares the backbone network and top-down path of existing lightweight baseline detection models (such as YOLOv5-S). To predict small-scale target regions in traffic scenes, an upsampling module and a prediction head composed of a sigmoid function are added to the baseline model's top-down path. The new module, based on the principle of flexibility and ease of construction, inherits the network structure from the baseline model's top-down path, such as... Figure 2 The lower half is shown.

[0032] To supervise the training of the foreground prediction module, the following loss function is used:

[0033]

[0034] In the formula, S is the mask representing the target region generated in S1, and S′ is the region representing the traffic target output by the prediction head of the foreground prediction module. and These represent the IOU loss and the binary cross-entropy loss, respectively. α is the weighting coefficient, which is set to 0.2 in this case.

[0035] The target detection module shares the bottom-up path and detection head of existing lightweight baseline models (such as YOLOv5-S) to characterize the category and location information of traffic targets. To specifically fuse traffic target features, a foreground target region weighted fusion method is established, such as... Figure 2 The SFGM module is shown in the diagram. Specifically, the target region predicted by the foreground prediction module is used to weight the encoded and decoded features of the foreground prediction module, as expressed below:

[0036]

[0037] In the formula, These represent channel concatenation operations, element-by-element addition, and element-by-element multiplication, respectively. This represents the operation consisting of convolution, batch normalization, and the SILU activation function. S′ is the region representing the traffic target output by the prediction head of the foreground prediction module. b and f d These represent the features corresponding to a 4x downsampling rate in the encoding and decoding paths of the foreground prediction module, respectively. β represents the learnable weighting coefficient. The fused F... sfIt highlights the expression of traffic target features, and through the layer-by-layer encoding of the target detection module, it can effectively extract traffic target features.

[0038] To supervise the learning of the student model, the following loss function was constructed:

[0039]

[0040] In the formula, For the loss in formula (1), This represents the loss of the baseline detection model. γ is the weighting coefficient, set to 0.15 in this case.

[0041] S3, based on the lightweight student model constructed in S2, establishes a teacher model TEOD-Net with strong feature representation capabilities. Specifically, the teacher model TEOD-Net and the student model STOD-Net have the same network structure, but the width of each layer of TEOD-Net is increased by 2.5 times, and the number of feature representation modules (such as CSP modules) in each network layer is increased by three times. The constructed teacher model has a large number of parameters, which can extract rich foreground region features and decode target type and location information. Consistent with the training of the student model, formula (3) is used to train the teacher model.

[0042] S4 establishes an effective distillation strategy, which mainly guides the foreground region feature interpretation and traffic target location and category information. Specifically, the rich traffic target region features decoded by the foreground prediction module of the teacher model constructed in S3 are used to guide the foreground prediction module of the student model constructed in S2 in the feature representation of the foreground region in the traffic scene. The distillation loss is as follows:

[0043]

[0044] In the formula, This represents the mean squared error loss. and These represent the features decoded by the foreground prediction module of the student model and the teacher model, respectively.

[0045] The features fused by the SFGM module in the teacher model are used to guide the weighted fusion of the encoding and decoding features of the foreground prediction module in the SFGM module of the student model. The distillation loss is as follows:

[0046]

[0047] In the formula, This represents the mean squared error loss. and These represent the features of the student model and the teacher model after fusion by the SFGM module, respectively.

[0048] The traffic target category and location information decoded by the target detection module of the teacher model are used to guide the target detection module of the student model in representing the target category and location in the traffic scene. The distillation loss is as follows:

[0049]

[0050] In the formula, This represents the mean squared error loss. and These represent the features decoded by the detection heads of the target detection modules in the student model and the teacher model, respectively.

[0051] S5. Based on formula (3), the traffic target detection training set in S1 and the generated foreground target region mask are used to train the teacher model. When the mAP value on the validation set reaches its maximum, the network weights of the teacher model are saved.

[0052] S6. The obtained teacher model network weights, the traffic target detection training set in S1, and the generated foreground target region mask are used to jointly supervise the training of the student model, thereby improving the student model's ability to represent traffic targets. The training loss is as follows:

[0053]

[0054] In the formula, The loss is defined by formula (3). The losses are defined by formulas (4), (5), and (6), respectively. ε is the weighting coefficient, which is set to 0.15 in this case.

[0055] S7. When the mAP value on the validation set reaches its maximum, save the student model parameters and use the data in the test set to test the detection performance of the lightweight student model.

[0056] Table 1 shows the quantitative detection results between different advanced models. It can be seen that the method proposed in this patent can effectively improve the detection capability of the baseline model. Specifically, the proposed method, through the design of an effective distillation strategy (foreground region feature guidance and traffic target location and category information guidance), transfers the strong feature representation capability of the teacher model to the student model, increasing only a slight amount of model parameters (0.1M-0.4M), yet improving the detection accuracy of GELAN-s, YOLOv5-s, and YOLOv7-tiny by 4.0%, 4.3%, and 4.6%, respectively. The outstanding detection capability demonstrates the effectiveness and generalization ability of the proposed method. The visualized detection results between the baseline model and the optimized model are shown in Table 1. Figure 3 As shown, the proposed method significantly improves the detection capability of the baseline model for multi-scale targets in complex traffic scenarios.

[0057] It is important to note that during the training phases of both the teacher and student models, the input image resolution was adjusted to 544×320 pixels, the batch size was set to 12, and data augmentation employed horizontal flipping and scaling strategies. The hyperparameters were consistent with the baseline detection model used. During the student model testing phase, the input image resolution was set to 544×320 pixels, the batch size was set to 1, and no data augmentation methods were used.

[0058] Table 1 Quantitative detection results of different models

[0059]

[0060] The above description is only a preferred embodiment of the present invention. It should be noted that for those skilled in the art, several improvements and modifications can be made without departing from the technical principles of the present invention, and these improvements and modifications should also be considered within the scope of protection of the present invention.

Claims

1. A lightweight traffic target detection method based on knowledge distillation, characterized in that, Includes the following steps: S1. Generate a foreground target region mask based on the traffic target detection training set; S2. Based on the lightweight detection model using the PANet network structure, a lightweight student model is established. The lightweight student model includes: a first foreground prediction module and a first target detection module. The first foreground prediction module shares the backbone network and top-down path of the lightweight detection model to predict the area where the target is located in traffic scenarios. The first target detection module shares the bottom-up path and detection head of the lightweight detection model, which is used to characterize the category and location information of traffic targets; S3. Based on S2, a lightweight student model is constructed, and a teacher model with strong feature expression capabilities is established. The teacher model includes: a second foreground prediction module and a second object detection module. The second foreground prediction module is used to decode rich features related to the traffic target area; The second target detection module is used to decode the category and location information of traffic targets; S4. Establish a distillation strategy, which includes: foreground area feature guidance and traffic target location and category information guidance; Specifically, the rich traffic target region features decoded by the foreground prediction module of the teacher model are used to guide the foreground prediction module of the student model to express the features of the foreground region in the traffic scene; the traffic target category and location information decoded by the target detection module of the teacher model are used to guide the target detection module of the student model to express the target category and location in the traffic scene. S5. Use the traffic target detection training set and the generated foreground target region mask from S1 to train the teacher model and obtain the network weights. S6. The obtained teacher model network weights, the traffic target detection training set in S1, and the generated foreground target region mask are used to jointly supervise the training of the student model, thereby improving the student model's ability to represent traffic targets. S7. Save the student model parameters and use the data in the test set to test the detection performance of the lightweight student model.

2. The lightweight traffic target detection method based on knowledge distillation as described in claim 1, characterized in that, In S1, the publicly available BDD100K dataset is used as the training set for model training, validation, and testing.

3. The lightweight traffic target detection method based on knowledge distillation as described in claim 2, characterized in that, In S1, 5K images and corresponding labels are randomly selected from the BDD100K training set for the validation set, and the remaining 65K images and corresponding labels are used for the new training set. The 10K images from the BDD100K validation set are used as the test set. When training the foreground prediction module, a foreground target region mask is generated based on the target detection labels corresponding to 65K training images. The process is as follows: set the region value within the bounding box in the target detection label to 1, and set the other region values ​​to 0.

4. The lightweight traffic target detection method based on knowledge distillation as described in claim 3, characterized in that, In S2, a lightweight student model STOD-Net is built based on the lightweight model YOLOV5-S. The first foreground prediction module shares the backbone network and top-down path of the existing lightweight baseline detection model. On the basis of the top-down path of the baseline detection model, an upsampling module and a prediction head composed of a Sigmoid function are added. To supervise the training of the foreground prediction module, the following loss function is used: In the formula, S is the mask representing the target region generated in S1, and S′ is the region representing the traffic target output by the prediction head of the foreground prediction module; and α represents IOU loss and binary cross-entropy loss, respectively; α is the weighting coefficient.

5. The lightweight traffic target detection method based on knowledge distillation as described in claim 4, characterized in that, In S2, the first target detection module shares the bottom-up path and detection head of the existing lightweight baseline model to characterize the category and location information of traffic targets. To fuse traffic target features, a foreground target region weighted fusion method is established, using the target region predicted by the foreground prediction module as the weighted encoding and decoding features of the foreground prediction module, as shown in the following expression: In the formula, These represent channel concatenation operations, element-wise addition, and element-wise multiplication, respectively. This represents the operation consisting of convolution, batch normalization, and the SILU activation function; S′ is the region representing the traffic target output by the prediction head of the foreground prediction module; f b and f d represents the features corresponding to a 4x downsampling rate in the encoding and decoding paths of the foreground prediction module, respectively; β represents the learnable weighting coefficient; and the fused F sf The expression of traffic target features is highlighted, and the traffic target features are extracted through layer-by-layer encoding by the target detection module; To supervise the learning of the student model, the following loss function is constructed: In the formula, For the loss in formula (1), γ represents the loss of the baseline detection model, and γ is the weighting coefficient.

6. The lightweight traffic target detection method based on knowledge distillation as described in claim 5, characterized in that, In S3, a teacher model TEOD-Net is established. The teacher model TEOD-Net and the student model STOD-Net have the same network structure. The width of each layer of TEOD-Net is increased by 2.5 times and the number of feature representation modules in each network layer is increased by 3 times. Formula (3) is used to train the teacher model.

7. The lightweight traffic target detection method based on knowledge distillation as described in claim 6, characterized in that, In S4, the rich traffic target region features decoded by the foreground prediction module of the teacher model built in S3 are used to guide the foreground prediction module of the student model built in S2 to represent the features of the foreground region in the traffic scene. The distillation loss is as follows: In the formula, This represents the mean squared error loss. and These represent the features decoded by the foreground prediction module of the student model and the teacher model, respectively. The features fused by the SFGM module in the teacher model are used to guide the weighted fusion of the encoding and decoding features of the foreground prediction module in the SFGM module of the student model. The distillation loss is as follows: In the formula, This represents the mean squared error loss. and These represent the features fused by the SFGM module in the student model and teacher model, respectively. The traffic target category and location information decoded by the target detection module of the teacher model are used to guide the target detection module of the student model in representing the target category and location in the traffic scene. The distillation loss is as follows: In the formula, This represents the mean squared error loss. and These represent the features decoded by the detection heads of the target detection modules in the student model and the teacher model, respectively.

8. The lightweight traffic target detection method based on knowledge distillation as described in claim 7, characterized in that, In S5, when the mAP value on the validation set reaches its maximum, the network weights of the teacher model are saved.

9. The lightweight traffic target detection method based on knowledge distillation as described in claim 8, characterized in that, In S6, the training loss is as follows: In the formula, The loss defined by formula (3), The losses are defined by formulas (4), (5), and (6), respectively; ε is the weighting coefficient.

10. The lightweight traffic target detection method based on knowledge distillation as described in claim 9, characterized in that, In S7, when the mAP value on the validation set reaches its maximum, the student model parameters are saved, and the data in the test set is used to test the detection performance of the lightweight student model.

Images