A target detection method, electronic device, and medium based on knowledge distillation

Abstract

This invention discloses a target detection method, electronic device, and medium based on knowledge distillation, comprising: inputting a target detection image to be detected into a pre-trained target detection model to obtain target detection results; wherein, the training process of the target detection model includes: acquiring target detection data and constructing it into a labeled dataset and an unlabeled dataset; training a teacher model using the balanced labeled dataset; inputting the unlabeled dataset into the teacher model to obtain confidence-labeled data; combining the unlabeled dataset and the confidence-labeled data, and after filtering, obtaining a confidence dataset; setting a distillation loss function, and training a student model based on the distillation loss function using the confidence dataset to obtain a target student model, i.e., a target detection model; wherein the distillation loss function is the sum of bounding box loss, target loss, and classification loss; wherein the target loss and classification loss are both constructed based on the confidence output by the target teacher model.

Description

Technical Field

[0001] This invention relates to the field of computer vision, specifically to a target detection method, electronic device, and medium based on knowledge distillation. Background Technology

[0002] Object detection is a core task in computer vision and digital image processing, widely used in autonomous driving, industrial inspection, security monitoring, and other fields. With the development of technologies such as convolutional neural networks, the accuracy of object detection models has continuously improved, but their computational requirements and deployment costs are also gradually increasing. Therefore, how to improve computational efficiency and reduce deployment costs while ensuring the accuracy of object detection models has become a research hotspot in the field of object detection in recent years.

[0003] Knowledge distillation (KD) is a general model compression method. Its basic principle is to transfer knowledge from a complex model (teacher model) to a simpler model (student model) using a specific distillation method, thereby obtaining a simpler model with higher performance and better results. There are two main methods for knowledge transfer between the teacher and student models: one is feature imitation, where the feature maps output by the intermediate networks of the student model mimic the feature maps output by the intermediate networks of the teacher model; the other is logit mimicking, where the final output of the student model mimics the final output of the teacher model. Existing knowledge distillation methods for object detection models mainly focus on feature imitation. This method relies heavily on a large amount of accurate labeled data in the training dataset for feature map selection and the determination of distillation regions on the feature maps. However, obtaining labeled data is difficult, time-consuming, and costly.

[0004] In the current field of object detection, the pursuit of higher detection accuracy has led to more complex object detection models, resulting in higher computational requirements and deployment costs. Therefore, how to implement a simple and easy-to-use knowledge distillation method in object detection tasks with less labeled data, effectively compressing the size of object detection models, improving object detection efficiency, and reducing model deployment costs, while maintaining high detection accuracy after compression, is an important requirement and a challenge. Summary of the Invention

[0005] To address the shortcomings of existing technologies, this invention provides a target detection method, electronic device, and medium based on knowledge distillation.

[0006] A first aspect of this invention provides a target detection method based on knowledge distillation, the method comprising:

[0007] The target detection image is input into a pre-trained target detection model to obtain the target detection result;

[0008] The training process of the object detection model includes:

[0009] Obtain object detection data; label a portion of the object detection data to construct labeled and unlabeled datasets;

[0010] Log-weighted class balancing is performed on the labeled dataset. The teacher model is then trained using the balanced labeled dataset until the teacher model converges, thus obtaining the target teacher model.

[0011] Input the unlabeled dataset into the target teacher model to obtain confidence-labeled data;

[0012] The unlabeled dataset is combined with the confidence-labeled data, and then the confidence-labeled dataset is constructed by performing class confidence distillation, log-weight class balancing, and target confidence filtering.

[0013] A distillation loss function is set, and a student model is trained based on the distillation loss function using a confidence dataset until the student model converges to obtain the target student model, which is the object detection model; wherein, the distillation loss function is the sum of the bounding box loss, the object loss, and the classification loss; the object loss and the classification loss are both constructed based on the confidence output of the target teacher model.

[0014] In a second aspect, an electronic device is provided, including a memory and a processor, wherein the memory is coupled to the processor; wherein the memory is used to store program data, and the processor is used to execute the program data to implement the above-described target detection method based on knowledge distillation.

[0015] A third aspect of the present invention provides a computer-readable storage medium having a computer program stored thereon, wherein the program, when executed by a processor, implements the above-described target detection method based on knowledge distillation.

[0016] A fourth aspect of the present invention provides a computer program product, including a computer program / instructions that, when executed by a processor, implement the above-described target detection method based on knowledge distillation.

[0017] Compared with the prior art, the beneficial effects of the present invention are as follows:

[0018] This invention provides a target detection method based on knowledge distillation. This method can reduce the dependence on labeled data and effectively reduce the cost of labeling target detection data. The confidence constraint method can effectively extract knowledge from the teacher model and transfer it to the student model, reducing the number of model parameters while maintaining high detection accuracy of the target detection model. The distillation loss function combined with confidence information can reduce the impact of low-quality labeled data on the student model during training, improve the detection accuracy of the student model, and improve the accuracy of target detection. Attached Figure Description

[0019] To more clearly illustrate the technical solutions in the embodiments of the present invention, the drawings used in the description of the embodiments will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For those skilled in the art, other drawings can be obtained based on these drawings without creative effort.

[0020] Figure 1 This is a schematic diagram of the target detection method based on knowledge distillation provided in an embodiment of the present invention;

[0021] Figure 2 This is an example diagram of automotive paint defects provided in an embodiment of the present invention;

[0022] Figure 3 This is a schematic diagram of an electronic device provided in an embodiment of the present invention. Detailed Implementation

[0023] To make the objectives and technical solutions of the embodiments of the present invention clearer, the technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some, not all, of the embodiments of the present invention. All other embodiments obtained by those skilled in the art based on the embodiments of the present invention without creative effort are within the scope of protection of the present invention.

[0024] It should be noted that, unless otherwise specified, the features in the following embodiments and implementation methods can be combined with each other.

[0025] like Figure 1 As shown, this embodiment of the invention provides a target detection method based on knowledge distillation, the method specifically including the following steps:

[0026] Step S1: Obtain object detection data; label a portion of the object detection data to construct labeled and unlabeled datasets.

[0027] Specifically, in this example, 20% to 50% of the target detection data is manually labeled to construct a labeled dataset, while the remaining data serves as an unlabeled dataset. It's important to note that in this example, the total number of target categories in the target detection data is greater than or equal to two.

[0028] The labeled dataset consists of Py image files containing the targets to be detected and Py text files storing the target annotation information. Each text file stores the annotation information of A targets contained in the corresponding image file. The annotation information is stored in A lines of text, with one line of text corresponding to the annotation information of one target. The annotation information stored in each line of text is arranged in the format [cls, x, y, w, h], where cls represents the category number of the target corresponding to the line of text. For the N categories in the labeled dataset 10, the category number cls is set to (0, 1, ..., N-1), x and y represent the pixel coordinates of the center point of the target bounding box on the image, and w and h represent the pixel width and pixel height of the target bounding box on the image.

[0029] The unlabeled dataset does not require manual annotation and consists only of Pn image files containing the target to be detected.

[0030] Step S2: Perform log-weight class balancing on the labeled dataset, and train the teacher model using the balanced labeled dataset until the teacher model converges to obtain the target teacher model.

[0031] Specifically, step S2 includes the following sub-steps:

[0032] Step S201: Calculate the percentage p of the number of targets in each category of the labeled dataset. i First, based on the category number cls in the annotation information, count the number of targets contained in each category. For N categories in the labeled dataset, the count of targets in each category is represented as (l1, l2, ..., l...). N The total number of all targets in the labeled dataset is The proportion of targets included in category i is p. i Represented as:

[0033]

[0034] Step S202, based on the target quantity percentage p i Calculate the sampling probability e corresponding to this category. i and according to probability e i Extract data from this category to perform data augmentation operations. The target quantity percentage is p. i The lower the category, the higher the probability of being drawn. iTherefore, a logarithmic weighting method is designed to calculate the sampling probability e. i The calculation formula is:

[0035]

[0036] Step S203, based on probability e i The extracted data undergoes data augmentation operations to achieve log-weighted class balance. These operations include eight types: rotation, scaling, cropping, horizontal flipping, vertical flipping, random noise reduction, brightness adjustment, and contrast adjustment. For each image file to be augmented and its corresponding annotation file, one of these eight operations is selected with probability P. E =0.5 Randomly select each data augmentation operation and combine them into a data augmentation operation set C E Then, perform corresponding data augmentation operations on the current image file and its corresponding annotation file. It is important to note that for each extracted image file and its corresponding annotation file, a new data augmentation operation set C is generated following the above steps. E Data augmentation operation set C E It includes probability P E =0.5 randomly selected k data augmentation operations, satisfying 1≤k≤8.

[0037] Step S203, set the data augmentation operation C E The training process begins by inputting the teacher model and continues until the teacher model converges, yielding the target teacher model.

[0038] It is important to note that the data augmentation operation set input to the teacher model is divided into a training set and a validation set. The training set accounts for 60% to 80% of the dataset input to the teacher model, and the remaining part of the dataset after removing the training set is used as the validation set.

[0039] Step S3: Input the unlabeled dataset into the target teacher model to obtain confidence-labeled data.

[0040] The confidence labeling format for the teacher model output is designed as follows: It is saved in a text file as multi-line text, with each line storing the confidence labeling information for one target, where C obj This represents the confidence level of the target inferred by the teacher model as indicating whether the target exists. This represents the category confidence level of the distilled target belonging to each category.

[0041] The unlabeled dataset is input into the teacher model for inference to obtain the inference results, i.e., the confidence-labeled data. The confidence-labeled data is then encoded in the format [cls, x, y, w, h, C]. obj ,(z1,...,z NThe format is saved in a text file.

[0042] Step S4: Combine the unlabeled dataset with the confidence-labeled data, and then perform class confidence distillation, log-weight class balancing, and target confidence filtering to construct the confidence dataset.

[0043] The category confidence distillation operation requires processing the probability (z1, ..., z2) of the target inferred by the teacher model to belong to each category in the teacher model's inference results (i.e., confidence-labeled data). N Distillation is performed to obtain class confidence levels. Category confidence distillation is represented as follows:

[0044]

[0045] Where T is a parameter representing the distillation temperature, and T > 0 is required.

[0046] The confidence label data after category confidence distillation is formatted as follows: The data is stored in a text file, and log-weight class balancing and target confidence filtering are performed to obtain the confidence dataset. It should be noted that the log-weight class balancing operation is the same as the log-weight class balancing operation described in step S2.

[0047] The target confidence filtering operation is based on the target confidence C in the confidence labeling data. obj Calculate the rejection probability p r And based on the elimination probability p r Randomly removing corresponding confidence level labels is a measure taken to reduce low-quality labeled data in the confidence dataset 50, thereby improving the training quality of the student model. The removal probability p... r The calculation method is p r =-log(C obj ), making the target confidence level C obj The lower the target value, the higher the elimination probability p. r When the target's elimination probability p r When the value is greater than 0.5, the elimination probability p is used as a guide. r Randomly delete the annotation information corresponding to the target when p r If the value is less than 0.5, the target will not be removed.

[0048] Step S5: Set the distillation loss function, and train the student model based on the distillation loss function using the confidence dataset until the student model converges to obtain the target student model, which is the target detection model; wherein, the distillation loss function is the sum of the bounding box loss, the target loss, and the classification loss; the target loss and the classification loss are both constructed based on the confidence output of the target teacher model.

[0049] To enable the student model to learn knowledge from the teacher model through knowledge distillation, a distillation loss function based on confidence constraints needs to be designed for the student model. During training, the student model learns data from B images in each batch, and the distillation loss L for that batch is expressed as:

[0050]

[0051] Among them, L box L represents the bounding box loss of the student model. obj L represents the target loss of the student model. cls Let B represent the classification loss of the student model, and let B represent the data of B images learned in a batch during the training of the student model, where b = 1, 2, 3, ..., B.

[0052] Bounding box loss L box Bounding boxes predicted by the student model [x s y s w s h s [x] and the bounding boxes inferred by teacher model 30 stored in the confidence dataset. t y t w t h t The Intersection over Union (IoU) is calculated to obtain the bounding box loss L for A targets in a batch. box The calculation formula is:

[0053]

[0054] The specific calculation method for the intersection-union ratio (IoU) is as follows:

[0055] x1 = max(x s -w s / 2,x t -w t / 2)

[0056] y1=min(y s +h s / 2,y t +h t / 2)

[0057] x2=min(x s +w s / 2,x t +w t / 2)

[0058] y2 = max(y s -h s / 2,y t -ht / 2)

[0059] Area I =|x1-x2|·|y1-y2|

[0060] Area U =w s h s +w t h t -Area I

[0061]

[0062] Target loss L obj The target confidence C is stored in the confidence dataset and is determined by the intersection-union ratio (IoU). obj The constraint calculation yields the following formula:

[0063]

[0064] l obj =-{iou·log[σ(s)} obj )]+(1-iou)·log[1-σ(s obj )]}·C obj

[0065] Where, σ(s) obj )=1 / [1+exp(-s obj )], s obj C represents the probability that the target predicted by the student model exists. obj A represents the confidence level of the target inferred by the teacher model, and A represents the number of targets predicted by the student model in an image, where a = 1, 2, 3, ..., A.

[0066] Classification loss L cls It combines the class confidence scores of the target belonging to N classes, stored in the confidence score dataset. The constraint calculation yields the following formula:

[0067]

[0068]

[0069] Where N represents the number of categories in the labeled dataset, A represents the number of targets predicted by the student model in an image, and a = 1, 2, 3, ..., A; t n This represents the target category encoding inferred by the teacher model (when the target is labeled as class n in the confidence dataset, t). n =1,t i≠n=0), s n The probability that the target predicted by student model 60 belongs to the nth class. The confidence level for the target inferred by the teacher model to belong to the nth class.

[0070] It should be noted that the classification loss L cls The calculation formula can also be:

[0071]

[0072]

[0073] The above two classification losses L cls The calculation formula can be selected based on actual performance.

[0074] The confidence dataset is input into the student model to begin the training process until the student model converges. It is important to note that the student model uses the entire confidence dataset as its training set, and the validation set used by the student model should be the same as the validation set used by the teacher model.

[0075] Step S6: The trained student model is used as the object detection model to perform the object detection task. The image to be detected is input into the student model for inference, and the object detection model outputs the object detection results. The detection results for each object are stored in a text file in the format [cls, x, y, w, h, conf], where cls is the category number of the object, x and y are the pixel coordinates of the center point of the object's bounding box in the image, w and h are the width and height of the object's bounding box, and conf is the probability that the student model predicts the object belongs to category cls.

[0076] It should be noted that, in this example, the specific network structures of the teacher and student models in this embodiment of the invention are not limited, as long as they can achieve object detection. The network structure of the teacher model is more complex than that of the student model. Existing object detection models, including but not limited to the R-CNN, YOLO, and SSD series, can be selected as the teacher model. Simpler object detection models, including but not limited to those from the R-CNN, YOLO, and SSD series, can be selected as the student model. Alternatively, based on existing object detection models, a new object detection model can be designed and improved according to the object detection task to be performed and the characteristics of the dataset, serving as the teacher model. A new, simpler model can be designed based on the teacher model, serving as the student model.

[0077] It should be noted that the embodiments of the present invention compress an untrained neural network model for object detection. The embodiments of the present invention do not compress each layer of the neural network model individually, but rather compress the entire neural network model.

[0078] The following comparative experiments, using the YOLOv7 series object detection model and a self-built automotive paint defect dataset, further illustrate the present invention.

[0079] The dataset used in the experiment is a dataset of automotive paint defects, including four types of defects: fiber, particles, sheet metal, and pinholes. Figure 2 As shown, the dataset is divided into three parts: a labeled dataset, a confidence dataset, and a test dataset. The labeled dataset, created manually, is used to train the teacher model and contains 2435 images of size 448×448 and their corresponding annotation files, divided into training and validation sets in a 7:3 ratio. The confidence dataset, obtained by inferring unlabeled data from the trained teacher model, is used to train the student model and contains 6225 images of size 448×448 and their corresponding annotation files. The confidence dataset is not divided and is used as the training set for the student model. The student model and the teacher model use the same validation set. The test dataset, created manually, is used to test the improvement effect of this patent on the detection accuracy of the student model and contains 398 images of size 448×448 and their corresponding annotation files.

[0080] The experiment used the YOLOv7 model as the teacher model and conducted comparative experiments on two student models, YOLOv7 and YOLOv7-Tiny, to test the technical effectiveness of this patent. The experimental group used a confidence dataset and trained the student model using the knowledge distillation method proposed in this patent; the control group also used a confidence dataset and trained the student model using conventional training methods, with the confidence information in the dataset not participating in the training process.

[0081] For training the student model, the initial learning rate was set to 0.001, and the SGD optimization strategy was used. With a momentum term of 0.9, a weight decay coefficient of 0.0005, and a batch size of 16, the model was trained for 200 epochs until the loss value of the student model stabilized, indicating that the student model was close to convergence. The mean accuracy (mAP@.5) at an IoU threshold of 0.5 was used to measure the detection accuracy of the student model.

[0082] Table 1: The effect of knowledge distillation on improving the accuracy of student model detection

[0083]

[0084] As shown in Table 1 above, compared with conventional training methods, the knowledge distillation method proposed in this invention significantly improves the detection accuracy of both YOLOv7 and YOLOv7-Tiny student models. Specifically, the detection accuracy of the YOLOv7 model is improved by 0.084, and the YOLOv7-Tiny model by 0.059. For small-area particle defects and sheet metal defects with blurred edges, manual annotation is difficult, making it challenging for conventional training methods to achieve high detection accuracy for these two types of defects. The knowledge distillation method proposed in this invention, based on confidence information, avoids the negative impact of low-quality manual annotation on the student model, significantly improving the detection accuracy of the student model for these two types of defects. Specifically, the detection accuracy for particle defects is improved by an average of 0.23, and the detection accuracy for sheet metal defects is improved by an average of 0.05. If YOLOv7-Tiny is used as the final student model, the number of parameters in the student model is about 1 / 6 of that in the teacher model YOLOv7. While achieving a good model compression effect, the detection accuracy is also improved by 0.006. This proves that the knowledge distillation method proposed in this invention can effectively compress the model while maintaining a level of detection accuracy comparable to that of the large model before compression.

[0085] Accordingly, this application also provides an electronic device, comprising: one or more processors; a memory for storing one or more programs; and, when the one or more programs are executed by the one or more processors, causing the one or more processors to implement the knowledge distillation-based target detection method as described above. Figure 3 The diagram shown is a hardware structure diagram of any device with data processing capabilities where the target detection method based on knowledge distillation provided in this embodiment of the invention is implemented, except... Figure 3 In addition to the processor, memory, and network interface shown, any data processing device in the embodiment may also include other hardware depending on the actual function of the data processing device, which will not be described in detail here.

[0086] Accordingly, this application also provides a computer-readable storage medium storing computer instructions thereon, which, when executed by a processor, implement the target detection method based on knowledge distillation as described above. The computer-readable storage medium can be an internal storage unit of any data-processing device as described in any of the foregoing embodiments, such as a hard disk or memory. The computer-readable storage medium can also be an external storage device, such as a plug-in hard disk, smart media card (SMC), SD card, flash card, etc., equipped on the device. Furthermore, the computer-readable storage medium can include both internal storage units of any data-processing device and external storage devices. The computer-readable storage medium is used to store the computer program and other programs and data required by the data-processing device, and can also be used to temporarily store data that has been output or will be output.

[0087] Other embodiments of this application will readily occur to those skilled in the art upon consideration of the specification and practice of the disclosure herein. This application is intended to cover any variations, uses, or adaptations of this application that follow the general principles of this application and include common knowledge or customary techniques in the art not disclosed herein. The specification and embodiments are to be considered exemplary only.

[0088] It should be understood that this application is not limited to the precise structure described above and shown in the accompanying drawings, and various modifications and changes can be made without departing from its scope.

Claims

1. A target detection method based on knowledge distillation, characterized in that, The method includes: The target detection image is input into a pre-trained target detection model to obtain the target detection result; The training process of the object detection model includes: Obtain object detection data; label a portion of the object detection data to construct labeled and unlabeled datasets; Log-weighted class balancing is performed on the labeled dataset. The teacher model is then trained using the balanced labeled dataset until the teacher model converges, thus obtaining the target teacher model. Input the unlabeled dataset into the target teacher model to obtain confidence-labeled data; The unlabeled dataset is combined with the confidence-labeled data, and then the confidence-labeled dataset is constructed by performing class confidence distillation, log-weight class balancing, and target confidence filtering. A distillation loss function is set, and a student model is trained based on the distillation loss function using a confidence dataset until the student model converges, thus obtaining the target student model, which is the object detection model; wherein, the distillation loss function is the sum of the bounding box loss, the object loss, and the classification loss; the object loss and the classification loss are both constructed based on the confidence output by the target teacher model; The process of performing log-weighted class balancing on labeled datasets includes: The labeled dataset has N categories, and the i-th category has 1 target, and calculate the percentage of targets in each category. The expression is: ； Percentage of target quantity The lower the category i, the higher the extraction probability should be. Extracting data from this category, extraction probability Represented as: ； Based on the sampling probability The extracted data is subjected to data augmentation operations to obtain a set of data augmentation operations. The process of combining unlabeled datasets with confidence-labeled data, and then performing class confidence distillation, log-weighted class balancing, and target confidence filtering to construct the confidence dataset includes the following steps: The category confidence score is obtained by distillation from the results of the teacher model inference; the confidence score for the i-th category is... Represented as: ； in, Let T represent the probability that the target obtained by the teacher model inference belongs to the i-th category, and let T represent the distillation temperature. For each objective, based on the objective confidence level Calculate the rejection probability When the probability of removing a target is greater than a threshold, the removal probability is determined based on the threshold. Randomly delete the annotation information corresponding to the target. If the rejection probability is less than the threshold, the target will not be rejected.

2. The target detection method based on knowledge distillation according to claim 1, characterized in that, The process of constructing a labeled dataset by annotating a portion of the target detection data includes: The annotations are stored in a multi-line text format, with each line storing the annotation information for one target. Where cls represents the category number of the current target, x and y represent the coordinates of the center point of the current target bounding box, and w and h represent the width and height of the current target bounding box.

3. The target detection method based on knowledge distillation according to claim 2, characterized in that, The process of inputting the unlabeled dataset into the target teacher model to obtain confidence-labeled data includes: The confidence level annotation data is stored in a multi-line text format, with each line storing the annotation information for one target. ,in The confidence score representing the existence of the target obtained from the teacher model's inference. This represents the class confidence level of the target belonging to each class after the class confidence distillation operation.

4. The target detection method based on knowledge distillation according to claim 1, characterized in that, The expression for the distillation loss function is as follows: ； In the formula, This represents the bounding box loss of the student model. This represents the target loss of the student model. Let B represent the classification loss of the student model, and let B represent the data of B images learned in a batch during the training of the student model, where b = 1, 2, 3, ..., B.

5. The target detection method based on knowledge distillation according to claim 4, characterized in that, The expression for the target loss of the student model is as follows: ； In the formula, , This represents the probability that the target predicted by the student model exists. A represents the confidence level of the target inferred by the teacher model, and A represents the number of targets predicted by the student model in an image, where a = 1, 2, 3, ..., A.

6. The target detection method based on knowledge distillation according to claim 4, characterized in that, The expression for the classification loss of the student model is: ； In the formula, N represents the number of categories in the labeled dataset, A represents the number of targets predicted by the student model in an image, and a = 1, 2, 3, ..., A; This represents the target category encoding inferred by the teacher model, when the target is labeled as class n in the confidence dataset. ; The confidence level for the student model to predict that the target belongs to the nth class. The confidence level for the target inferred by the teacher model to belong to the nth class; or, 。 7. An electronic device comprising a memory and a processor, characterized in that, The memory is coupled to the processor; wherein the memory is used to store program data, and the processor is used to execute the program data to implement the target detection method based on knowledge distillation as described in any one of claims 1-6.

8. A computer-readable storage medium having a computer program stored thereon, characterized in that, When the program is executed by the processor, it implements the target detection method based on knowledge distillation as described in any one of claims 1-6.

Images