Method for training an object detection model

Abstract

The present invention provides a method for training an object detection model. The method includes the following steps: selecting a teacher model and a student model through a computing device, freezing the parameters of the entire teacher model and the parameters after layer N+1 of the head network of the student model in the first stage of training, training the teacher model and the student model so that the student model gradually approaches the output of the teacher model, obtaining a first-stage student model, and training the parameters of the head network of the first-stage student model in the second stage to improve the accuracy of the first-stage student model.

Description

[Technical Field]

[0001] This invention relates to an object detection model training method, and more particularly to an object detection model training method for improving the accuracy of small, lightweight, and fast-running object detection models and compressing the models, so that the trained models are more suitable for applications in edge computing devices. [Previous Technology]

[0002] Fine-tuning is a common technique in machine learning and deep learning. It mainly refers to further training an existing pre-trained model on a small dataset for a specific task to adapt it to the requirements of a new task. It can quickly achieve efficient model training while saving computing resources and data.

[0003] Object detection is a core task in the field of computer vision, which aims to identify all target objects in an image or video and simultaneously determine their classification and location.

[0004] Compared to image classification, which only labels whether a certain type of object exists in an image, object detection needs to provide more refined results, including the object's bounding box and its category. From the perspective of network or model structure, it can be divided into backbone network, neck network, and head network.

[0005] The purpose of the backbone is feature extraction, extracting low- to high-level feature representations from the image; the purpose of the neck is feature fusion enhancement and multi-scale processing, allowing objects of different sizes in the image to be accurately detected simultaneously; the head is the task layer or output layer, whose purpose is to complete object classification and bounding box regression. The head network for object detection is usually the main fine-tuning part of downstream computer vision tasks such as object detection, while the backbone and neck are basically used to train the model's ability to extract and fuse features from the image, and do not directly involve too many object detection-related capabilities.

[0006] In addition, due to the rise of edge computing devices, manufacturers are researching how to put lighter but more accurate models into edge devices. However, these lightweight models are often less accurate than large models. Therefore, if lightweight models can learn or imitate large models, they can achieve higher accuracy while maintaining lightweight design.

[0007] Therefore, in order to achieve the above objectives, the industry uses knowledge distillation techniques. Knowledge distillation is mainly used to enable a relatively inaccurate student model to learn a more accurate feature representation from a more accurate teacher model while being trained with ground truth. However, general knowledge distillation methods have some limitations: (1) Object detection is a specific downstream task of computer vision, and distillation techniques require special handling. (2) A large amount of labeled data is still required, but in real-world applications, data is not easy to obtain, and the labeling and cleaning of data is quite time-consuming. (3) In addition, traditional methods often require the simultaneous use of ground truth and teacher model predictions to train the student model. However, if only the output of the teacher model is used as the simulated label, then ground truth is not required, and it can become "unsupervised". However, doing so alone is often insufficient to enable the student model to obtain sufficient accuracy in practice.

[0008] Therefore, in order to overcome the limitations of knowledge distillation and achieve better accuracy, this invention can freeze the parameters other than the head network during training, and obtain excellent results by only fine-tuning the head network. In addition to saving a lot of training time, it can also effectively improve accuracy; therefore, this invention should be the best solution. [Summary of the Invention]

[0009] The object detection model training method of the present invention comprises the following steps: (1) Selecting a teacher model and a student model through a computing device, wherein the teacher model and the student model are both object detection models (e.g., YOLO models), and the object detection models (e.g., YOLO models) of the teacher model and the student model are different; (2) Performing a first-stage training of the student model through the computing device, wherein the student model is trained using random initial parameters, and during the first-stage training process, freezing the parameters of the entire teacher model and freezing the parameters of the N+1th layer and beyond of the head network of the student model, the frozen parameters will not be updated during the training process; (3) Then, training is performed using a first dataset, and the difference between the outputs of the teacher model and the student model is calculated based on the outputs of the teacher model and the student model, and the training is repeated a certain number of times so that the student model gradually approaches the output of the teacher model, and finally a first-stage student model can be obtained; (4) The first-stage student model is loaded into the computing device, and a second dataset with labeled data is used to train the parameters of the head network of the first-stage student model in the second stage. After repeated training a certain number of times, the accuracy of the first-stage student model is improved, thereby updating the parameters of the head network of the first-stage student model, and finally a second-stage student model can be obtained.

[0010] More specifically, the first dataset is either a dataset without labeled data or a dataset with labeled data.

[0011] More specifically, the evaluation of the difference between the output of the teacher model and the output of the student model is carried out by calculating a loss function.

[0012] More specifically, the student model gradually approximates the output of the teacher model, so that the student model learns the intermediate feature representation (latent space) and domain knowledge of the teacher model.

[0013] More specifically, in the first stage of training, only the parameters of the Nth layer and earlier layers of the head network of the student model can be trained and updated; while in the second stage of training, the parameters of the head network of the student model are fine-tuned and updated.

[0014] More specifically, before the first stage of training, the Nth layer of the head network of the student model is connected to the N+1th layer of the head network of the teacher model so that the student model can learn the feature distribution of the intermediate feature representation of the teacher model.

[0015] More specifically, freezing the parameters of the entire teacher model means that the parameters of the teacher model will not be updated during the first stage of training, so as to maintain the original state.

[0016] More specifically, freezing the parameters of the head network of the student model after layer N+1 means that during the first stage of training, the parameters of the head network of the student model before layer N can be trained and updated.

[0017] More specifically, the second-stage training of the parameters of the head network of the first-stage student model means that during the second-stage training process, only the parameters of the head network of the student model are trained and updated, and the parts other than the head network are frozen.

Implementation Method

[0018] Other technical contents, features and effects of the present invention will be clearly presented in the following detailed description of the preferred embodiments with reference to the accompanying drawings.

[0019] The connection described below in this invention refers to the connection behavior generated through physical wired or physical wireless or logical connection methods, so as to enable modules performing specific functions of electronic components to generate unidirectional or / and bidirectional transmission of digital signals and / or analog signals and / or digitized / discrete analog-analog signals.

[0020] As used herein, the articles “a,” “one,” and “any” refer to the grammar of one or more (i.e., at least one) items. For example, “an element” means one element or more.

[0021] As used herein, “unit”, “device”, “module” and “equipment” in this case refer to any hardware name that can achieve the specific functions and / or purposes of this case, but the implementation of this case is not limited to hardware with such names.

[0022] The term "backbone network" as used below refers to the backbone network.

[0023] The term "neck network" as used below refers to the neck network.

[0024] The term “head network” as used below in this case refers to the head network.

[0025] Please refer to Figure 1, which is a flowchart of the object detection model training method of the present invention. The steps are as follows: (1) Select a teacher model and a student model through a computing device, wherein the teacher model and the student model are both object detection models (e.g., YOLO models), and the object detection models (e.g., YOLO models) of the teacher model and the student model are different 101; (2) Perform the first stage training of the student model through the computing device. The student model is trained using random starting parameters. During the first stage training, the parameters of the entire teacher model and the parameters of the N+1th layer of the head network of the student model are frozen. The output of the Nth layer of the student network is connected in series to the N+1th layer of the teacher network as its input. The frozen parameters will not be updated during the training process 102; (3) Then, a first dataset is used for training, and the difference between the outputs of the teacher model and the student model is calculated based on the outputs of the teacher model and the student model. The training is repeated a certain number of times so that the student model gradually approaches the output of the teacher model. Finally, a first-stage student model 103 can be obtained. (4) The first-stage student model is loaded through the computing device, and a second dataset with labeled data is used to train the parameters of the head network of the first-stage student model in the second stage. After repeated training a certain number of times, the accuracy of the first-stage student model is improved, thereby updating the parameters of the head network of the first-stage student model. Finally, a second-stage student model 104 can be obtained.

[0026] The computing device 1 is any computing device capable of computing and execution control, such as a desktop computer, a laptop computer, or a server device. As shown in Figure 2, the computing device 1 includes at least one processor 11 and at least one computer-readable recording medium 12, wherein the computer-readable recording medium 12 contains one or more applications 121, and the computer-readable recording medium 12 further stores computer-readable instructions. When the processor 11 executes the computer-readable instructions, the application 121 can be made to run.

[0027] Step 103, the first dataset mentioned is either a dataset without labeled data or a dataset with labeled data.

[0028] Step 103, wherein the difference between the output of the teacher model and the output of the student model is evaluated by calculating the loss function based on the output of the teacher model and the output of the student model; at the same time, the student model gradually approaches the output of the teacher model, which is used to enable the student model to learn the intermediate feature representation (latent space) and domain knowledge of the teacher model.

[0029] Step 102, during the first stage of training, the parameters of the entire teacher model and the parameters of the N+1th layer and beyond of the head network of the student model are frozen; wherein freezing the parameters of the entire teacher model means that the parameters of the teacher model are not updated during the first stage of training, so as to maintain the original state; and the parameters of the N+1th layer and beyond of the head network of the student model means that during the first stage of training, only all parameters of the student model other than those frozen are updated.

[0030] Step 102 aims to train only all parameters except those of the student network's frozen layer, and to facilitate the calculation of its output and ensure effective convergence of the model update, the output of the Nth layer of the student network's head network is connected in series to the N+1th layer of the teacher network's head network.

[0031] In addition, before the first stage of training, the Nth layer of the head network of the student model is connected to the N+1th layer of the head network of the teacher model so that the student model can learn the feature distribution of the intermediate feature representation of the teacher model.

[0032] Step 104, wherein the parameters of the head network of the first-stage student model are trained in the second stage, means that in the second-stage training process, only the parameters of the head network of the student model are trained and updated, and the parts other than the head network are frozen; in this way, the parameters of the head network of the student model can be finely adjusted and updated in the second-stage training process.

[0033] To better understand the technology of this case, the features of this case are described in more detail below. Regarding the technology of knowledge distillation, the core of this case is very similar to fine-tuning. However, if only the prediction value of the teacher model is used to simulate the real ground truth for the student model to learn, then the ground truth (labeled data) is not needed, thus achieving unsupervised learning. However, simply using this method will inevitably fail to successfully train the student model. Therefore, this case splits the original knowledge distillation into two parts, as follows: (1) The first stage of training adopts unsupervised learning, allowing the student model to learn the output and representation of the teacher model. The purpose is to make the latent space output by the student model as close as possible to the latent space of the teacher model. (2) The second stage of training adopts fine-tuning, allowing the student model that learns the teacher model to achieve higher and usable accuracy through ground truth fine-tuning. (3) Combining the above two stages, the method adopted in this case is different from general knowledge distillation. The method in this case is multi-stage and has the meaning of breaking down the characteristics of knowledge distillation and reversing the operation. Therefore, it is defined as reverse distillation (Inverted KD - Knowledge Distillation).

[0034] The following provides a more detailed implementation description of the technology in this case. However, the model types mentioned in the description are for implementation purposes only, and this case is not limited to a single model.

[0035] The first stage of training in this case is described as follows: (1) A sufficiently accurate pre-trained model is selected as the teacher model. In this invention, YOLOv10M (trained using the COCO dataset) is selected. (2) A sufficiently lightweight and fast inference target network is selected as the student model. In this invention, YOLOv10N is selected. (3) A large amount of dataset related to the target task is collected, and only a portion is selected for annotation, approximately 1 / 4 of the total. In this invention, two publicly available datasets are used for this simulation experiment. The datasets are as follows: (a) The first dataset is COCO, with approximately 118,287 training images in 80 categories. (b) The second dataset is COCO-mintrain, which is about 1 / 4 of the dataset taken from COCO. For example, 25,000 images are extracted from COCO, which can be considered a "mini version". (c) In order to save annotation costs, this case only annotates about 1 / 4 of the images, and the other 3 / 4 of the images are not annotated and are used for unsupervised training; the definition of 1 / 4 in this case is only a certain implementation mode, and its proportion can be adjusted as needed.

[0036] The first stage of training in this case, as shown in Figure 3, is described as follows: (1) In the training process of the first stage, the student model 2 (YOLOv10N) uses random initial parameters instead of pre-trained parameters. (2) In this case, CrossKD is first used for optimization. The forward inference of the "Nth layer of the head network" of the student model 2 (YOLOv10N) is transferred to the "N+1th layer of the head network" of the teacher model 3 (YOLOv10M). N is the second layer of the head network. That is, in the structure of the YOLO head, a specified layer is selected, its output features are spliced ​​or aligned, and compared with the output of the corresponding layer of the teacher model 3. The purpose is to let the student model 2 (YOLOv10N) learn the domain knowledge and latent space of the teacher model 3 (YOLOv10M) in this stage, and can be trained using unlabeled data. (3) During the first phase of training, the parameters of the head network of the entire teacher model 3 (YOLOv10M) and student model 2 (YOLOv10N) after layer N+1 are frozen; the parameters of the entire teacher model 3 will not be updated during training, and only the student model 2 (layer N and earlier layers) can be trained and updated. (4) Calculate the unsupervised distillation loss as follows: (a) Input the image into the teacher model 3 (YOLOv10M) and obtain the output of the teacher model 3 (P t). (b) Input the same image into the student model 2 (YOLOv10N) and obtain the output of the student model 22 (P s). (c) Calculate the loss function 𝐿𝑜𝑠𝑠𝑐𝑘𝑑 to measure the difference between the output of the student model and the output of the teacher model 3. (5) Since all parameters after layer N+1 of the head network of the overall teacher model 3 (YOLOv10M) and student model 2 (YOLOv10N) are frozen, backpropagation will not update the frozen model parameters. The frozen parameters after layer N+1 of student model 2 (YOLOv10N) will be fine-tuned in the second stage of training. (6) Remove the labeled data from COCO or COCO-mintrain (train using the dataset with unlabeled data) and perform the above training for 60 epochs.

[0037] The second stage of training in this case, as shown in Figure 4, is implemented as follows: (1) The parameters and structure of the student model 4 (YOLOv10N) trained in the first stage are extracted and used as the initial model for the second stage training. (2) COCO-mintrain (all data are labeled) is used as the training set, and the parameters of the backbone network and neck network of the student model 4 (YOLOv10N) are frozen. Only the head network is trained. The reasons are as follows: (a) In the structure of YOLO, the backbone network and neck network are mainly responsible for feature extraction and fusion. The student model 2 (YOLOv10N) has already learned the ability of the teacher model 3 (YOLOv10M) for feature extraction and fusion, as well as some object detection in the first stage of unsupervised knowledge distillation. In this stage, only a smaller amount of data is needed to fine-tune the head network of the student model 4 (YOLOv10N). (b) If the overall parameters of student model 4 (YOLOv10N) are fine-tuned together, student model 4 may gradually forget the knowledge learned from teacher model 3. Therefore, the second stage of training only trains the head network. (3) As shown in Figure 3, the image is input into the network, and after forward propagation, the output Ps of student model 4 is obtained. Then, the loss function Loss org is calculated with the ground truth Pgt. The model parameters are updated by backpropagation using the gradient obtained from the loss function. (4) After 60 epochs of training, the loss function and accuracy will tend to stabilize and saturate. At this time, the training of the target network is completed and can be tested and used.

[0038] This case further compares the trained model with existing models, and the comparison is illustrated below: (1) The models used for training and comparison in this case are: (a) Target model (student model): YOLOv10N (b) Teacher model (teacher model): YOLOv10M (c) Comparison benchmark model: YOLOv10N. The benchmark model used for comparison in this case is obtained by training YOLOv10N with random initial model parameters for 60 epochs using COCO-mintrain, which is the same hyperparameter used in the training of this invention. (d) The parameter sizes (YOLOv10N, YOLOv10S, YOLOv10M) and inference times of different models are compared as shown in Table 1. Table 1: Size and Inference Time of Different Models Model Test Size #Params FLOPs Latency YOLOv10N 640 2.3M 6.7G 1.84ms YOLOv10S 640 7.2M 21.6G 2.49ms YOLOv10M 640 15.4M 59.1G 4.74ms (e) To demonstrate that using a larger amount of data for unsupervised distillation in the first stage can improve accuracy, the dataset for unsupervised knowledge distillation in the first stage was changed from COCO-mintrain to COCO. Of course, using COCO-mintrain for unsupervised knowledge distillation in the first stage would also achieve better accuracy compared to other models or traditional knowledge distillation methods. (f) In Table 2 below, InvertedKD refers to this invention. The first minus sign indicates the training dataset used in the first stage of training, and the second minus sign indicates the layer to which the data is frozen. Table 2: Comparison Results of Different Models Method Teacher Student mAP50 mAP50-95 Baseline (YOLOv10N) - YOLOv10N 0.2915 0.193 CrossKD YOLOv10M YOLOv10N 0.332 0.221 InvertedKD-COCOmintrain YOLOv10M YOLOv10N 0.368 0.249 Inverted KD-COCO-22 YOLOv10M YOLOv10N 0.474 0.326 (g1) Wherein, InvertedKD-COCOmintrain refers to the training dataset used in the first stage of training of this invention being COCOmintrain. (g2) Wherein, InvertedKD-COCO-22 refers to the training dataset used in the first stage of training of this invention being COCO, and frozen to the 22nd layer of the network. (g3) Wherein, mAP (Mean Average Precision): In object detection, mAP is a key indicator for measuring the accuracy of model prediction, mainly used to evaluate the relationship between "precision" and "recall" of object detection. (g4) Wherein, the mAP metric has different calculation methods, the main difference being the IoU threshold (Thresholds, page 6 of 14 (Invention Specification)). (g5) Wherein, mAP@50([email protected]) means that the IoU is fixed at 0.5 (i.e., IoU ≥ 50% is considered True Positive), and the [email protected] of all categories is calculated and then averaged. (g6) Wherein, mAP@50-95 (COCO mAP, mAP@[0.5:0.95]) represents IoU from 0.5 to 0.95, with intervals of 0.05, for a total of 10 IoU thresholds. The AP of each IoU threshold is calculated and then averaged. Generally, mAP@50 is usually higher (because the IoU requirement is lower), and mAP@50-95 is usually lower (because a higher IoU requirement is needed to be considered True Positive). (g7) Wherein, CrossKD is a knowledge distillation (KD) method. Its main goal is to distill the feature knowledge at different levels (between layers), rather than simply guiding student network learning through the logits of the final output layer. This invention optimizes this method into an unsupervised distillation method and uses this optimized method in the first stage of the training step in this invention. (g) The following conclusions can be drawn from Table 2: (h1) Using InvertedKD can eliminate the conflict in model parameter updates that may be caused by different target loss functions during simultaneous training. (h2) The unsupervised distillation method in the first paragraph can be selected for specific tasks. For example, in this invention, CrossKD is changed from supervised distillation to unsupervised distillation, and the loss function is changed from Kullback-Leibler Divergence to Binary Cross Entropy.Because CrossKD achieves relatively superior results in knowledge distillation for object detection, it is optimized to achieve unsupervised learning and multi-stage distillation. (h3) Since the first stage uses unsupervised distillation, a large amount of labeled data is not required; only the dataset needed for fine-tuning in the second stage is needed. (h4) Using InvertedKD yields higher accuracy than direct training (baseline-YOLOv10N) and other knowledge distillation methods (such as CrossKD). Even if the first and second stages use the same dataset, as long as the dataset is large enough, superior results can be obtained. (h5) As shown in Table 2 above, InvertedKD can achieve better results by using different knowledge distillation methods in the first stage and different fine-tuning methods in the second stage. (h6) In Table 2 above, the InvertedKD of this invention uses an unsupervised optimized version of CrossKD in the first stage and only fine-tunes the detection head in the second stage, which not only saves training resources and time but also achieves superior accuracy.

[0039] The object detection model training method provided by the present invention has the following advantages when compared with other conventional technologies: (1) In order to overcome the limitations of knowledge distillation and achieve better accuracy, the present invention can first use a large amount of unlabeled data to allow the student model to learn the "feature extraction ability" from the teacher model, and then use a small amount of labeled data for refinement; this can significantly reduce the labeling cost and achieve better performance. (2) In the first stage of training, the present invention only needs to use the output of the teacher model and can train the student model without ground truth, and then by comparing the intermediate feature representation (latent space) and the final output of the two, the student model can initially learn the domain knowledge of the teacher model. (3) In the second stage of training, the present invention only needs to use a small amount of real labeled data to fine-tune the student model for similar teacher models, so that its detection results are consistent with the actual target task and improve accuracy.

[0040] The present invention has been disclosed above through the above embodiments, but it is not intended to limit the present invention. Any person skilled in the art can make some modifications and refinements after understanding the foregoing technical features and embodiments of the present invention, without departing from the spirit and scope of the present invention. Therefore, the scope of patent protection of the present invention shall be determined by the claims attached to this specification. [Simplified Explanation of the Diagram]

[0041] [Figure 1] is a flowchart illustrating the object detection model training method of the present invention. [Figure 2] is a schematic diagram illustrating the architecture of the computing device in the object detection model training method of the present invention. [Figure 3] is a schematic diagram illustrating the first-stage training architecture of the object detection model training method of the present invention. [Figure 4] is a schematic diagram illustrating the second-stage training architecture of the object detection model training method of the present invention.

Claims

1. A method for training an object detection model, comprising the following steps: Selecting a teacher model and a student model through a computing device, wherein both the teacher model and the student model are object detection models, but the object detection models of the teacher model and the student model are different; Performing a first-stage training of the student model through the computing device, wherein the computing device uses random initial parameters to train the student model, and during the first-stage training process, freezing all parameters of the teacher model and freezing the parameters of the N+1th layer and beyond of the head network of the student model, the frozen parameters are not updated during the training process; The computing device uses a first dataset for training, and calculates the difference between the outputs of the teacher model and the student model based on the outputs of the two models, and then repeatedly trains a certain number of times so that the student model gradually approaches the output of the teacher model, and finally obtains a first-stage student model; The computing device loads the first-stage student model and uses a second dataset with labeled data to train the parameters of the head network of the first-stage student model in the second stage. After repeated training a certain number of times, the accuracy of the first-stage student model is improved, thereby updating the parameters of the head network of the first-stage student model. Finally, a second-stage student model is obtained. Before the first stage of training, the Nth layer of the student model's head network is connected to the N+1th layer of the teacher model's head network. This allows the student model to learn the feature distribution of the teacher model's intermediate feature representations. Freezing the parameters after the N+1th layer of the student model's head network means that during the first stage of training, the parameters before the Nth layer of the student model's head network, as well as the parameters of the backbone and neck networks, can be trained and updated. During the first stage of training, only the Nth layer, the layers before the Nth layer, and all parameters of the backbone and neck networks of the student model's head network can be trained and updated. In the second stage of training, the parameters of the student model's head network are fine-tuned and updated. The second stage of training on the parameters of the student model's head network from the first stage means that during the second stage of training, only the parameters of the student model's head network are trained and updated, while other parts of the network are frozen.

2. The object detection model training method as described in claim 1, wherein the first dataset is either a dataset without labeled data or a dataset with labeled data.

3. The object detection model training method as described in claim 1, wherein the difference between the output of the teacher model and the output of the student model is evaluated by calculating a loss function based on the output of the teacher model and the output of the student model.

4. The object detection model training method as described in claim 1, wherein the student model gradually approximates the output of the teacher model, so that the student model learns the intermediate feature representations and domain knowledge of the teacher model.

5. The object detection model training method as described in Request 1, wherein freezing the parameters of the entire teacher model means that the parameters of the teacher model are not updated during the first stage of training, so as to maintain the original state.

Images