Knowledge distillation-based target detection model training method

Abstract

The invention relates to a training method of a target detection model based on knowledge distillation, and belongs to the field of target detection. Comprising the following steps: constructing a training data set and a target detection model, and training the target detection model based on the training data set to output an inference graph; calculating a first loss function to obtain a trained target detection model, and taking the trained target detection model as a teacher model; simplifying parameters of the teacher model to obtain a student model; obtaining a second training data set based on the training data set, training the teacher model and the student model based on the second training data set, and respectively outputting prediction labels; and calculating a second loss function to obtain a trained student model as a trained target detection model. According to the method, the optimal distillation temperature and loss coefficient are obtained by adopting the differential mutation algorithm, and the trained target detection model is obtained by training the student model through knowledge distillation, so that the problem of low detection precision caused by the fact that the existing target detection model is trained by a small amount of samples is solved.

Description

Technical Field

[0001] The present invention relates to the field of target detection, and in particular to a training method for a target detection model based on knowledge distillation. Background Art

[0002] As an important branch of computer vision, object detection technology is widely used in security, unmanned driving, smart home and other fields. With the rapid development of artificial intelligence, object detection methods based on deep learning have become a research hotspot.

[0003] However, existing deep learning-based object detection models usually require a large number of samples for training. When trained with a small number of samples, the detection accuracy is usually low. In addition, existing deep learning-based object detection models are usually not simplified, which is not conducive to model deployment under conditions of limited hardware resources.

[0004] It can be seen that providing a method for training target detection models with a small number of samples is an urgent problem to be solved. Summary of the invention

[0005] In view of the above analysis, the present invention aims to provide a training method for a target detection model based on knowledge distillation, so as to solve the problem that the existing target detection model has low detection accuracy when trained with a small number of samples.

[0006] The present invention provides a method for training a target detection model based on knowledge distillation, the method comprising the following steps:

[0007] Preprocess the images in the public dataset and construct a training dataset after annotation; construct a target detection model, input each sample in the training dataset into the target detection model for training, and output the corresponding reasoning graph; the reasoning graph includes a prediction box;

[0008] The first loss function is calculated by the error between the labeled box of the sample and the predicted box of the corresponding reasoning graph. When the first loss function reaches a threshold, a trained target detection model is obtained and used as a teacher model. The number of convolution channels of the backbone network of the teacher model is halved to obtain a student model.

[0009] A second training data set is obtained based on the training data set, and each sample in the second training data set is input into the teacher model and the student model for training, and the teacher model and the student model respectively output corresponding prediction labels; a second loss function is calculated by the error between the true label of the sample and the predicted label corresponding to the student model and the error between the predicted label corresponding to the teacher model and the predicted label corresponding to the student model, and when the second loss function reaches a threshold, a trained student model is obtained, and the trained student model is used as a trained target detection model.

[0010] Furthermore, an initial target detection model is constructed based on the YOLOv5s model, wherein the initial target detection model includes an input end, a backbone network, a neck network and a detection end; a One-to-one Detect network structure is obtained based on the One-to-many Detect network structure of the detection end, and a One-to-one Detect network structure is added to the detection end to construct a target detection model.

[0011] Furthermore, the first loss function includes a loss function of a One-to-many Detect network structure and a loss function of a One-to-one Detect network structure; the loss function of the One-to-many Detect network structure includes a detection target confidence loss, a detection target category loss, and a prediction box loss; the loss function of the One-to-one Detect network structure includes a detection target category loss and a prediction box loss.

[0012] Furthermore, the second loss function includes a loss function of a student model and a distillation loss function; the loss function of the student model represents the error between the true label of the sample and the predicted label corresponding to the student model, and the distillation loss function represents the error between the predicted label corresponding to the teacher model and the predicted label corresponding to the student model.

[0013] Furthermore, the expression of the second loss function is as follows:

[0014] Loss total2 =αLoss student +βLoss distillation ,

[0015] Among them, Loss total2 Represents the second loss function, Loss student Represents the loss function of the student model, Loss distillation represents the distillation loss function, and α and β are weight coefficients.

[0016] Furthermore, the loss function expression of the student model is as follows:

[0017] Loss student =BCE(s, label)

[0018]

[0019] Among them, s is the true label of the sample, label is the predicted label corresponding to the student model, BCE(s, label) is the cross entropy loss between s and label, and z i The predicted label belonging to the i-th category output by the student model.

[0020] Furthermore, the distillation loss function expression is as follows:

[0021] Loss distillation =KL(q t ,q s )

[0022]

[0023] Among them, qs is the predicted label output by the teacher model, q t is the predicted label output by the student model, KL(q t ,q s ) is q s With q t The KL divergence of z is is the predicted label belonging to the i-th category output by the teacher model, z it is the predicted label belonging to the i-th category output by the student model, and T is the distillation temperature.

[0024] Furthermore, the parameters T, α, and β are obtained through the differential mutation algorithm:

[0025] Step S301, obtain constraint conditions according to the training accuracy of the student model, and randomly generate T, α, β generating individuals that meet the constraint conditions to construct an initial population;

[0026] Step S302: Calculate the loss function corresponding to each individual in the population as the fitness value of each individual, and select individuals whose fitness values ​​are less than a threshold to form an excellent population;

[0027] Step S303, determine whether the termination condition is met. If so, obtain parameters T, α, and β based on the individual with the smallest fitness value in the population; otherwise, cross and mutate the individuals in the excellent population according to the crossover probability and mutation probability to obtain new individuals, combine the individuals in the excellent population with the new individuals to obtain an updated population, use the updated population as the population for the next iteration, and repeat steps S302 to S303.

[0028] Furthermore, the expression of the constraint condition is as follows:

[0029]

[0030] Among them, L is the second loss function; τ is the student model; w is the weight parameter of the student model.

[0031] Furthermore, the expression of the first loss function is as follows:

[0032] L total1 =L o2m +L o2o

[0033] L o2m =λ obj BCE obj +λ cls BCE cls +λ box EIU

[0034] L o2o =δ11BCE cls +δ2IoU,

[0035] Among them, L total1 represents the first loss function, L o2m Represents the loss function of the One-to-many Detect network structure, L o2o Represents the loss function of the One-to-one Detect network structure, BCE obj To detect the target confidence loss, BCE cls is the target category loss, EIoU is the prediction box loss of the One-to-many Detect network structure, IoU is the prediction box loss of the One-to-one Detect network structure, and λ obj , cls , box , δ1, δ2 are weight coefficients.

[0036] Compared with the prior art, the present invention can achieve at least one of the following beneficial effects:

[0037] 1. The present invention obtains a trained target detection model by training the student model through knowledge distillation. The accuracy of the lightweight model can be restored when only a small number of samples are used, which is convenient for deploying the model under conditions of limited hardware resources. It solves the problem that the existing target detection model has low detection accuracy when trained with a small number of samples.

[0038] 2. The present invention adopts a differential mutation algorithm to obtain the optimal distillation temperature, student model loss coefficient α, and distillation loss coefficient β, thereby improving the prediction accuracy of the student model.

[0039] 3. The present invention improves the detection end. Through the One-to-one Detect network structure, the prediction result can be directly output when used without post-processing the inference graph, thereby producing an end-to-end target detection effect.

[0040] 4. The C3AM module proposed in the present invention adds an attention mechanism to prune the convolutional channels to be fused through channel information enhancement, eliminates redundant channels in the network, reduces the parameters of the model, and improves the efficiency of model training and reasoning.

[0041] In the present invention, the above-mentioned technical solutions can also be combined with each other to achieve more preferred combination solutions. Other features and advantages of the present invention will be described in the subsequent description, and some advantages can become obvious from the description, or can be understood by practicing the present invention. The purpose and other advantages of the present invention can be realized and obtained through the contents particularly pointed out in the description and the drawings. BRIEF DESCRIPTION OF THE DRAWINGS

[0042] The drawings are only for the purpose of illustrating particular embodiments and are not to be considered limiting of the present invention. Like reference symbols denote like components throughout the drawings.

[0043] Figure 1 A flowchart of a method for training a target detection model based on knowledge distillation according to an embodiment of the present invention;

[0044] FIG2( a ) is a schematic diagram of the overall network structure of the initial target detection model according to an embodiment of the present invention;

[0045] FIG2( b ) is a schematic diagram of the network structure of the initial target detection model CBS, C3_1, C3_2, BottleNeck1, and SPPF modules according to an embodiment of the present invention;

[0046] Figure 3 This is a schematic diagram of the structure of the C3AM module according to an embodiment of the present invention;

[0047] Figure 4 It is a schematic diagram of the structure of the improved neck network according to an embodiment of the present invention;

[0048] Figure 5 Schematic diagram of the network structure of the improved target detection model according to the embodiment of the present invention. DETAILED DESCRIPTION

[0049] The preferred embodiments of the present invention are described in detail below in conjunction with the accompanying drawings, wherein the accompanying drawings constitute a part of this application and are used together with the embodiments of the present invention to illustrate the principles of the present invention, but are not used to limit the scope of the present invention.

[0050] A specific embodiment of the present invention discloses a training method for a target detection model based on knowledge distillation. Figure 1 As shown, the method comprises the following steps:

[0051] Step S1: preprocess the images in the public data set, and construct a training data set after annotation; construct a target detection model, input each sample in the training data set into the target detection model for training, and output the corresponding reasoning graph; the reasoning graph includes a prediction box;

[0052] Step S2, calculating a first loss function by the error between the labeled box of the sample and the predicted box of the corresponding reasoning graph, and obtaining a trained target detection model when the first loss function reaches a threshold, and using it as a teacher model; reducing the number of convolution channels of the backbone network of the teacher model by half to obtain a student model;

[0053] Step S3, obtaining a second training data set based on the training data set, inputting each sample in the second training data set into the teacher model and the student model for training, and the teacher model and the student model respectively output corresponding prediction labels; calculating a second loss function through the error between the true label of the sample and the predicted label corresponding to the student model and the error between the predicted label corresponding to the teacher model and the predicted label corresponding to the student model, and when the second loss function reaches a threshold, obtaining a trained student model, and using the trained student model as a trained target detection model.

[0054] Specifically, in step S1, the preprocessing of images in the public dataset and the construction of the training dataset after annotation include:

[0055] Get a public dataset and randomly select some images from it;

[0056] Preprocess the selected images by removing the mean, converting image channels, and converting image formats;

[0057] Adaptively scale the preprocessed image to a uniform size;

[0058] Perform image enhancement processing on the adaptively scaled image;

[0059] The adaptively scaled and enhanced images are annotated to construct a training dataset.

[0060] Specifically, obtain a public dataset, such as the COCO dataset, the visdrone dataset, and the vehicles dataset, and randomly select some images from them. The de-meaning refers to subtracting the pixel values ​​of the three channels R, G, and B of the selected image from the mean of the corresponding channels of the ImageNet dataset. The image channel conversion refers to converting the channel format of the selected image to the RGB format. The image format conversion refers to converting the format of the selected image to the same format, such as the jpeg format.

[0061] It can be understood that removing the mean of the selected image can reduce the deviation of the pixel distribution of the input image.

[0062] Furthermore, the adaptively scaling the preprocessed image to a uniform size includes:

[0063] Set the target size; for each preprocessed image:

[0064] The larger one is selected from the width and height of the image, and the ratio of the width or height corresponding to the target size to the selected width or height of the image is used as the scaling ratio of the image;

[0065] scaling the image according to the scaling ratio of the image;

[0066] If the size of the image after scaling is different from the target size, the image after scaling is padded with a fixed pixel value based on the difference between the size of the image after scaling and the target size, so that the size of the image after filling is consistent with the target size and the central content remains unchanged.

[0067] For example, the target size is set to 640×640. Assuming that the size of an image is 1280×720, the width of the image is selected as 1280, and the ratio of the width of the target size 640 to 1280 is As the image zooms in, press The image is reduced, and the size of the image after reduction is 640×360. According to the difference between 640×360 and 640×640, the image is padded with pixel values ​​of (114, 114, 114) around it after reduction, so that the size of the image after filling is 640×640 and the central content remains unchanged.

[0068] It should be noted that the preprocessed image is scaled and filled with fixed pixel values ​​around the scaled image, so that the preprocessed image is adaptively scaled to a uniform size and the content of the original image is centered.

[0069] Furthermore, the image enhancement processing on the adaptively scaled image includes one or more of the following methods:

[0070] Randomly select part of the content in the image and copy it to other images to generate a new image;

[0071] Performing a two-dimensional linear transformation on the image to generate a new image; the two-dimensional linear transformation includes rotation, translation, scaling, shearing, and flipping;

[0072] Randomly select two images from the images, mix the two images in a certain ratio, and generate a new image;

[0073] Randomly adjust the hue, saturation, and brightness of the image to generate a new image;

[0074] Part of the image is randomly blocked to generate a new image.

[0075] It should be noted that the shearing includes horizontal shearing and vertical shearing. Horizontal shearing achieves a distortion effect by adding a horizontal displacement on each row of the image, and vertical shearing achieves a distortion effect by adding a vertical displacement on each column of the image.

[0076] It can be understood that through image enhancement, not only the number of samples is increased, but also the background of the image is enriched, so that the target detection model can learn more diverse knowledge during training, providing a basis for improving the robustness and generalization ability of the model.

[0077] Furthermore, a dataset is constructed by annotating each target to be detected in the adaptively scaled and image enhanced image with a rectangular box through the LabelImg annotation tool; the target to be detected is selected as needed; the annotation content includes the category number, the coordinates of the center point of the target to be detected, and the height and width of the annotation box.

[0078] Specifically, the training data set includes adaptively scaled and enhanced images and their corresponding annotation contents.

[0079] It should be noted that the annotation content of each image is generated as a file in txt format with the same name as the image. Therefore, the training data set includes a set of images and a set of their corresponding annotation content files.

[0080] Furthermore, an initial target detection model is constructed based on the YOLOv5s model, and the initial target detection model includes an input end, a backbone network, a neck network and a detection end; the C3_2 module of the neck network is improved to obtain a C3AM module; based on the One-to-many Detect network structure of the detection end, a One-to-one Detect network structure is obtained; the C3_2 module in the neck network is replaced by the C3AM module, and a One-to-one Detect network structure is added to the detection end, thereby constructing a target detection model.

[0081] Specifically, the initial target detection model is constructed based on the YOLOv5s model, and its network structure diagram is shown in Figure 2(a) and Figure 2(b). The input end is used to obtain the image to be detected; the backbone network is used to extract feature information of multiple different scales from the image to obtain feature maps of multiple corresponding scales; the neck network includes a feature pyramid network (FPN) and a path aggregation network (PAN), which are used to fuse feature maps of multiple corresponding scales, fuse the shallow feature information of the image with the deep feature information, and input them to the detection end; the detection end is used to output target detection results of multiple corresponding scales. Among them, the CBS module includes a convolutional layer (Conv2d), a batch normalization layer (BN), and a SiLu activation function connected in sequence, which are used to extract feature information of multiple different scales of the image. The C3_1 module includes a first convolution module, a first feature fusion layer (Concat), and a first CBS module connected in sequence, and is used to fuse feature information of multiple different scales; wherein, the first convolution module includes a second CBS module and a first bottleneck module connected in parallel, the first bottleneck module includes a third CBS module and a first bottleneck layer (BottleNeck1) connected in sequence, the first bottleneck layer includes a residual module and a second feature fusion layer connected in sequence, the residual module includes a jump connection and a dual CBS module connected in parallel, and the dual CBS module includes a fourth and a fifth CBS module connected in sequence. The C3_2 module includes a second convolution module, a first feature fusion layer, and a first CBS module connected in sequence, and is used to fuse feature information of multiple different scales; the second convolution module includes a second CBS module and a second bottleneck module connected in parallel, the second bottleneck module includes a third CBS module and a second bottleneck layer (BottleNeck2) connected in sequence, and the second bottleneck layer includes a fourth and a fifth CBS module connected in sequence. The SPPF module includes a first CBS module, three maximum pooling layers (MaxPool2d), a feature fusion layer, and a second CBS module connected in sequence, and is used to extract global feature information of the image; wherein the output of the first CBS module and the output of the three maximum pooling layers are also jump-connected to the feature fusion layer.

[0082] Furthermore, if Figure 3As shown, the C3AM module includes a Shuffle module, a split layer (Split), the first to fourth CBS modules, a feature fusion layer, and an AM module; the input feature map is divided into two paths through the sequentially connected Shuffle modules and split layers, and the number of channels in each path is reduced to half of the original number of channels, one of which is input to the feature fusion layer after convolution operation through the first CBS module, and the other is input to the feature fusion layer after convolution operation through the second CBS module, and the output of the second CBS module is also input to the feature fusion layer after convolution operation through the third CBS module, and the output of the third CBS module is also input to the feature fusion layer after convolution operation through the fourth CBS module, and the feature fusion layer fuses the above inputs to obtain a fused feature map, and the fused feature map is input to the AM module for attention calculation to obtain a feature map with enhanced channel information.

[0083] Specifically, the Shuffle module is used to fuse spatial attention and channel attention, and the AM module is an attention mechanism module, which is used to obtain a feature map with enhanced channel information.

[0084] Furthermore, the fused feature map is input into the AM module for attention calculation to obtain the channel attention weighted feature map, which includes:

[0085] Step S101: Input the fused feature map of each channel into the AM module for pHash calculation to obtain the corresponding hash value; if the Hamming distance between the hash values ​​corresponding to the feature maps of two channels is less than the threshold, randomly retain one of the channels and remove the other channel;

[0086] Step S102, calculating the maximum eigenvalue and the average eigenvalue of the feature map of each channel after the channel is removed, generating a global maximum eigenvector based on the maximum eigenvalue, and generating an average eigenvector based on the average eigenvalue;

[0087] Step S103, performing convolution operation on the global maximum eigenvector and the average eigenvector with a 1×1 convolution kernel respectively, so as to learn the attention weight of each channel after channel removal; adding the global maximum eigenvector and the average eigenvector after the convolution operation to obtain a sum vector;

[0088] Step S104: Generate a channel attention weight vector based on the summed vector through a Sigmoid activation function;

[0089] Step S105, sort the channel attention weight values ​​in the channel attention weight vector from high to low, set the pruning ratio, set the lowest part of the channel attention weight value to 0 according to the pruning ratio to obtain an updated channel attention weight vector; multiply each channel attention weight value in the updated channel attention weight vector by the feature map of its corresponding channel to obtain a feature map after channel attention weighting.

[0090] For example, assume that the feature map after the fusion of the i-th channel is F i ∈R C×H×W , where C is the number of channels, H is the height of the fused feature map, and W is the width of the fused feature map; perform pHash calculation on the fused feature map of each channel to obtain the corresponding hash value; if the Hamming distance between the hash values ​​corresponding to the fused feature maps of two channels is less than 5, then randomly retain one of the channels and remove the other channel, and the feature map of the fused i-th channel after the channel removal is F′ i ∈R C′×H×W , where C′ is the number of channels after channel removal, C′≤C; calculate the maximum eigenvalue and average eigenvalue of the feature map of each channel after channel removal, and generate the global maximum eigenvector e based on the maximum eigenvalue of the feature map of each channel max ∈R C′×1×1 , generates the average feature vector e based on the average eigenvalue of the feature map of each channel min ∈R C′×1×1 , where e max =MaxPool(F′), e min =AvgPool(F′); the global maximum eigenvector e max and the average eigenvector e min Convolution operations are performed with 1×1 convolution kernels respectively to obtain e max ′=Conv 1×1 (e max ) and e min ′=Conv 1×1 (e min ), e max ′ and e min ' add, and we get the sum vector V = e max ′+e min ′,V∈R C′×1×1 ; Generate channel attention weight vector based on V through Sigmoid activation function The channel attention weight values ​​in the channel attention weight vector are sorted from high to low, and the pruning ratio is set to 40%, and the lowest 40% of the channel attention weight values ​​are set to 0 to obtain an updated channel attention weight vector; each channel attention weight value in the updated channel attention weight vector is multiplied by the feature map of its corresponding channel to obtain a feature map after channel attention weighting.

[0091] It can be understood that if the Hamming distance between the hash values ​​corresponding to the feature maps of two channels is less than the threshold, the feature maps of the two channels are considered similar, so one of the channels is randomly retained and the other channel is removed to reduce the number of channels. The global maximum eigenvector contains the maximum eigenvalue of the feature map of each channel after the channel is removed, and the average eigenvector contains the average eigenvalue of the feature map of each channel after the channel is removed. When the attention weight value of a channel is low, it is set to 0 according to the pruning ratio. At this time, the channel attention weight value multiplied by the feature map of its corresponding channel is also 0. Therefore, the channel corresponding to it will be suppressed, so that the channel information with greater correlation is enhanced and the channel information with less correlation is suppressed.

[0092] Furthermore, the One-to-oneDetect network structure obtained based on the One-to-many Detect network structure of the detection end includes:

[0093] Based on the training data set, the annotation boxes are clustered to obtain multiple anchor boxes corresponding to large, medium and small objects respectively;

[0094] The One-to-many Detect network structure of the detection end includes multiple convolutional layers, which are used to predict large, medium and small targets using corresponding multiple anchor frames to obtain multiple prediction frames;

[0095] A One-to-oneDetect network structure with the same structure is constructed based on the One-to-many Detect network structure of the detection end; the One-to-oneDetect network structure is used to select the anchor frame closest to the corresponding target from the corresponding multiple anchor frames for large, medium and small targets respectively, and each target obtains a predicted frame.

[0096] Furthermore, clustering the annotated boxes based on the training data set to obtain multiple anchor boxes corresponding to large, medium and small targets respectively includes: obtaining the sizes of all the annotated boxes of the training data set, that is, the width and height of the annotated boxes, using a clustering algorithm to perform cluster analysis on the sizes of all the annotated boxes, and obtaining anchor boxes corresponding to large, medium and small targets respectively; wherein each target corresponds to multiple anchor boxes of different sizes.

[0097] Exemplarily, a K-means clustering algorithm is used to perform cluster analysis on the sizes of all the annotation boxes, and three anchor boxes of different sizes corresponding to large, medium and small targets are obtained respectively. Therefore, there are a total of nine anchor boxes of different sizes corresponding to large, medium and small targets.

[0098] Furthermore, the One-to-many Detect network structure of the detection end is used to predict large, medium and small targets using corresponding multiple anchor frames, and each target obtains multiple prediction frames.

[0099] Exemplarily, the One-to-many Detect network structure of the detection end makes predictions on feature maps of sizes 80×80, 40×40, and 20×20, respectively. For example, three anchor boxes corresponding to large targets are used for prediction on a feature map of size 80×80, and each large target can obtain three corresponding prediction boxes.

[0100] It should be noted that the reasoning graph includes prediction boxes, each of which includes the coordinates of the center point corresponding to the prediction box, the height and width of the prediction box, the detection target category number, and the confidence score.

[0101] It is understandable that when prediction is performed through the One-to-many Detect network structure, the resulting inference graph contains a large number of target prediction boxes, many of which are redundant. Therefore, it is necessary to remove redundant prediction boxes through the NMS (non-maximum suppression) method, that is, the inference graph needs to be post-processed to obtain the predicted target.

[0102] Furthermore, selecting the anchor frame closest to the corresponding target from the multiple anchor frames corresponding to the large, medium and small targets respectively includes: calculating the aspect ratio of each anchor frame corresponding to the corresponding target, calculating the aspect ratio of the corresponding target annotation frame, and selecting the anchor frame closest to the aspect ratio of the corresponding target from the multiple anchor frames.

[0103] Furthermore, the One-to-one Detect network structure obtains the predicted box of each target through the following formula:

[0104] b x =σ(t x )×2-0.5+c x

[0105] b y =σ(t y )×2-0.5+c y

[0106] b w =p w (σ(t w )×2) 2

[0107] b h =p h (σ(t h )×2) 2 ,

[0108] Among them, b x is the x coordinate of the center point of the prediction box, b y is the y coordinate of the center point of the prediction box, b w is the width of the prediction box, b h is the height of the prediction box, t x is the x-coordinate of the center point of the prediction box output by the model, t y is the y coordinate of the center point of the prediction box output by the model, t w is the width of the prediction box output by the model, t h is the height of the prediction box output by the model, c x is the x coordinate of the upper left corner of the grid unit corresponding to the center point of the target annotation box, c y is the y coordinate of the upper left corner of the grid unit corresponding to the center point of the target annotation box, p w is the width of the closest anchor box selected, p h is the height of the closest anchor box selected.

[0109] Furthermore, if Figure 4 As shown, all C32 modules in the neck network of the initial target detection model are replaced with C3AM modules; and a parallel One-to-oneDetect network structure is added to the detection end to construct a target detection model. The network structure diagram of the target detection model is shown in Figure 5 shown.

[0110] It is understandable that the number of convolution channels in the initial target detection model is usually determined manually based on experience, and there is redundancy, which seriously affects the training and reasoning efficiency of the model, especially when used under conditions of limited hardware resources. By adopting the C3AM module, the convolution channels to be fused are pruned, the model parameters are reduced, and the efficiency of model training and reasoning is improved. The present invention uses a One-to-one Detect network structure to directly output prediction results when used without post-processing the reasoning graph, thereby producing an end-to-end target detection effect, while ensuring the model prediction accuracy. Improve the efficiency of reasoning; solve the problem that the target detection model used in the existing target detection method requires post-processing of the detection results, resulting in low detection speed.

[0111] Furthermore, each sample in the training data set is input into the target detection model for training, and the corresponding reasoning graph output includes: a One-to-many Detect network structure and a One-to-one Detect network structure outputting corresponding reasoning graphs respectively.

[0112] Specifically, in step S2, the first loss function includes the loss function of the One-to-many Detect network structure and the loss function of the One-to-one Detect network structure; the loss function of the One-to-many Detect network structure includes detection target confidence loss, detection target category loss, and prediction box loss; the loss function of the One-to-oneDetect network structure includes detection target category loss and prediction box loss; the expression of the first loss function is as follows:

[0113] L total1 =L o2m +L o2o

[0114] L o2m =λ obj BCE obj +λ cls BCE cls +λ box EIU

[0115] L o2o =δ11BCE cls +δ2IoU,

[0116] Among them, L total1 represents the first loss function, L o2m Represents the loss function of the One-to-many Detect network structure, L o2o Represents the loss function of the One-to-one Detect network structure, BCE obj To detect the target confidence loss, BCE cls is the target category loss, EIoU is the prediction box loss of the One-to-many Detect network structure, IoU is the prediction box loss of the One-to-one Detect network structure, and λ obj , cls , box , δ1, δ2 are weight coefficients.

[0117] Preferably, λ obj Set to 1.0, λ cls Set to 0.5, λ box Set it to 0.05, α to 0.5, and β to 0.1.

[0118] Furthermore, the detection target confidence loss represents the error in prediction accuracy, that is, whether there is a target in the prediction box. The loss function expression of the detection target confidence loss is as follows:

[0119]

[0120] Among them, N is the total number of samples, m is the number of positive samples, Nm is the number of negative samples, and y i =1 means there is a target in the prediction box, y i =0 means that there is no target in the prediction box, p1 is the probability that there is a target in the prediction box, and α is the coefficient.

[0121] It can be understood that negative samples correspond to the image background, and positive samples correspond to the actual target area in the image. Usually, there are more negative samples and fewer positive samples, which easily leads to the loss value of negative samples overwhelming the loss value of positive samples, thus reducing the convergence speed of the model and the accuracy of target detection. obj The loss function can dynamically adjust the loss values ​​of positive and negative samples, thereby improving the speed of model convergence and the accuracy of target detection.

[0122] Furthermore, the detection target category loss represents the error between the category of the prediction box and the category of the annotation box. The loss function expression of the detection target category loss is as follows:

[0123]

[0124] Where N p is the number of positive samples, n is the number of categories, C i =1 means the predicted category belongs to C i Category, C i =0 means the predicted category does not belong to C i Category, p i Indicates that the prediction is C i The probability of the class.

[0125] Furthermore, the prediction box loss represents the position error between the prediction box and the sample annotation box. The loss function expression of the prediction box loss of the One-to-many Detect network structure is as follows:

[0126]

[0127] Among them, b pre represents the area occupied by the prediction box, b gt represents the area occupied by the annotation box, IoU is the intersection-over-union ratio of the predicted box and the annotation box, ρ 2 (bpre , b gt ) represents the distance between the center point of the prediction box and the center point of the annotation box, w pre is the width of the prediction box, w gt is the width of the annotation box, ρ 2 (w pre , w gt ) represents the square of the difference between the width of the prediction box and the width of the annotation box, h pre is the height of the prediction box, h gt is the height of the annotation box, ρ 2 (h pre ,h gt ) represents the square of the height difference between the predicted box and the labeled box, c w 2 Represents the square of the minimum enclosing rectangle width of the prediction box and the annotation box, c h 2 Represents the square of the minimum enclosing rectangle height of the prediction box and the annotation box.

[0128] It should be noted that a positive sample refers to a sample whose error between the predicted box in the inference graph and the actual labeled box is less than a threshold, indicating that the intersection-over-union ratio of the predicted box and the labeled box is greater than the threshold and the category of the predicted box is consistent with the category of the labeled box, otherwise it is a negative sample. The prediction box loss of the One-to-one Detect network structure is the intersection-over-union ratio of the predicted box and the labeled box.

[0129] Specifically, the number of training rounds is set to 200. In each round of training, the batch size is set to 8, the optimizer is Adam, the initial learning rate is 0.01, the weight decay rate is 0.001, the momentum factor is 0.937, the weight decay factor is 0.0005, the detection target confidence loss threshold is 0.5, the detection target category loss threshold is 0.75, and the prediction box loss threshold is 0.5.

[0130] Specifically, in step S3, one third of the samples are extracted from the training data set to obtain a second training data set.

[0131] Furthermore, the second loss function includes a loss function of the student model and a distillation loss function; the loss function of the student model represents the error between the true label of the sample and the predicted label corresponding to the student model, and the distillation loss function represents the error between the predicted label corresponding to the teacher model and the predicted label corresponding to the student model; the expression of the second loss function is as follows:

[0132] Loss total2 =αLoss student +βLosd distillation ,

[0133] Among them, Loss total2 Represents the second loss function, Loss student Represents the loss function of the student model, Loss distillation represents the distillation loss function, and α and β are weight coefficients.

[0134] Furthermore, the loss function expression of the student model is as follows:

[0135] Loss student =BCE(s, label)

[0136]

[0137] Among them, s is the true label of the sample, label is the predicted label corresponding to the student model, BCE(s, label) is the cross entropy loss between s and label, and z i The predicted label belonging to the i-th category output by the student model.

[0138] Furthermore, the distillation loss function expression is as follows:

[0139] Loss distillation =KL(q t ,q s )

[0140]

[0141]

[0142] Among them, qs is the predicted label output by the teacher model, q t is the predicted label output by the student model, KL(q t ,q s ) is q s With q t The KL divergence of z is is the predicted label belonging to the i-th category output by the teacher model, z it is the predicted label belonging to the i-th category output by the student model, and T is the distillation temperature.

[0143] Furthermore, the parameters T, α, and β are obtained through the differential mutation algorithm:

[0144] Step S301, obtain constraint conditions according to the training accuracy of the student model, and randomly generate T, α, β generating individuals that meet the constraint conditions to construct an initial population;

[0145] Step S302: Calculate the loss function corresponding to each individual in the population as the fitness value of each individual, and select individuals whose fitness values ​​are less than a threshold to form an excellent population;

[0146] Step S303, determine whether the termination condition is met. If so, obtain parameters T, α, and β based on the individual with the smallest fitness value in the population; otherwise, cross and mutate the individuals in the excellent population according to the crossover probability and mutation probability to obtain new individuals, combine the individuals in the excellent population with the new individuals to obtain an updated population, use the updated population as the population for the next iteration, and repeat steps S302 to S303.

[0147] Specifically, in step S301, the expression of the constraint condition is as follows:

[0148]

[0149] Among them, L is the second loss function; τ is the student model; w is the weight parameter of the student model.

[0150] Specifically, T, α, and β that meet the constraints are randomly generated using the following formula:

[0151] T ij =9*rand ij (0, 1)+1

[0152] α ij =rand ij (0, 1)

[0153] β ij =rand ij (0, 1)

[0154] Among them, rand(0,1) represents a random decimal between [0,1], i represents the sequence number of the individual in the first population, and j represents the length of the individual's chromosome.

[0155] Specifically, each individual is represented by X = [T, α, β], and the initial population is constructed by randomly generating individuals that meet the constraints within the value range of T, α, and β according to the above formula.

[0156] Specifically, in step S302, the loss function Loss corresponding to each individual in the population is calculated. total2 As the fitness value of each individual.

[0157] Specifically, in step S303, the termination condition includes that the loss function of the individual with the smallest fitness value in the current population is less than a threshold or the number of iterations reaches a maximum value.

[0158] Furthermore, if the termination condition is not met, individuals are randomly selected from the excellent population according to the mutation probability and new individuals are generated using the following formula:

[0159]

[0160] Among them, F represents the mutation probability; p1, p2, and p3 are random integers, representing the sequence number of the individual in the excellent population; t represents the current number of iterations, and h ij (t+1) represents a newly generated individual;

[0161] Specifically, if the termination condition is not met, individuals are randomly selected from the excellent population according to the crossover probability and new individuals are generated using the following formula:

[0162]

[0163] Among them, randl ij is a random decimal between [0,1], CR is the crossover probability, x ij (t) is an individual randomly selected from the excellent population.

[0164] Specifically, individuals in the excellent population are combined with new individuals generated by crossover and mutation to obtain an updated population.

[0165] Furthermore, the image to be detected is input into the trained target detection model, the One-to-manyDetect network structure is frozen, the predicted label is directly output through the One-to-one Detect network structure, and it is determined whether the predicted label is the target to be detected, thereby realizing target recognition.

[0166] It is understandable that the selection of three key parameters, namely distillation temperature, distillation loss coefficient and student model loss coefficient, is usually determined manually based on experience, which seriously affects the learning effect of the actual student model. The present invention adopts a differential mutation algorithm to obtain the optimal distillation temperature, student model loss coefficient α, and distillation loss coefficient β, thereby improving the prediction accuracy of the student model. The present invention trains the student model through knowledge distillation to obtain a trained target detection model, which can restore the accuracy of the lightweight model when only a small number of samples are used, making it easy to deploy the model under conditions of limited hardware resources; it solves the problem of low detection accuracy of the existing target detection model when trained with a small number of samples.

[0167] Performance evaluation:

[0168] Specifically, performance evaluation indicators are set to evaluate the performance of the improved target detection model, including mean average precision (mAP), precision (P), and recall (R). Precision P is used to evaluate the accuracy of target detection, and recall R is used to evaluate the comprehensiveness of target detection. mAP refers to the average area under the PR curve of all categories, which comprehensively considers the precision and recall of the model.

[0169] Specifically, the calculation formulas for average accuracy, precision, and recall are as follows:

[0170]

[0171] Among them, TP represents the number of correct target detection results, FP represents the number of incorrect target detection results, FN represents the number of missed detections where the target was not detected, AP represents the area under the PR curve of the category, and n c is the total number of categories.

[0172] Specifically, the COCO dataset is used to train YOLOv5s and the target detection model improved by the method proposed in the present invention, wherein the parameters of the improved target detection model are 40% of the parameters of the YOLOv5s model. Each model is trained for 200 rounds, and the results of the above indicators are obtained as shown in Table 1:

[0173] Table 1 Experimental results

[0174]

[0175] It can be seen from the above table that compared with YOLOv5s, the target detection model improved by the method proposed in the present invention greatly reduces the number of parameters while also improving the average accuracy to a certain extent, effectively improving the detection speed.

[0176] Compared with the prior art, the training method of the target detection model based on knowledge distillation provided by the present invention has the following beneficial effects:

[0177] 1. The present invention obtains a trained target detection model by training the student model through knowledge distillation. The accuracy of the lightweight model can be restored when only a small number of samples are used, which is convenient for deploying the model under conditions of limited hardware resources. It solves the problem that the existing target detection model has low detection accuracy when trained with a small number of samples.

[0178] 2. The present invention adopts a differential mutation algorithm to obtain the optimal distillation temperature, student model loss coefficient α, and distillation loss coefficient β, thereby improving the prediction accuracy of the student model.

[0179] 3. The present invention improves the detection end. Through the One-to-one Detect network structure, the prediction result can be directly output when used without post-processing the inference graph, thereby producing an end-to-end target detection effect.

[0180] 4. The C3AM module proposed in the present invention adds an attention mechanism to prune the convolutional channels to be fused through channel information enhancement, eliminates redundant channels in the network, reduces the parameters of the model, and improves the efficiency of model training and reasoning.

[0181] Those skilled in the art will appreciate that all or part of the processes of the above-mentioned embodiments can be implemented by instructing related hardware through a computer program, and the program can be stored in a computer-readable storage medium, wherein the computer-readable storage medium is a disk, an optical disk, a read-only storage memory, or a random access memory, etc.

[0182] The above description is only a preferred specific implementation manner of the present invention, but the protection scope of the present invention is not limited thereto. Any changes or substitutions that can be easily conceived by any technician familiar with the technical field within the technical scope disclosed by the present invention should be covered within the protection scope of the present invention.

Claims

1. A training method for an object detection model based on knowledge distillation, characterized in that: The method comprises the following steps: Preprocess the images in the public dataset and construct a training dataset after annotation; construct a target detection model, input each sample in the training dataset into the target detection model for training, and output the corresponding reasoning graph; the reasoning graph includes a prediction box; The first loss function is calculated by the error between the labeled box of the sample and the predicted box of the corresponding reasoning graph. When the first loss function reaches a threshold, a trained target detection model is obtained and used as a teacher model. The number of convolution channels of the backbone network of the teacher model is halved to obtain a student model. A second training data set is obtained based on the training data set, and each sample in the second training data set is input into the teacher model and the student model for training, and the teacher model and the student model respectively output corresponding prediction labels; a second loss function is calculated by the error between the true label of the sample and the predicted label corresponding to the student model and the error between the predicted label corresponding to the teacher model and the predicted label corresponding to the student model, and when the second loss function reaches a threshold, a trained student model is obtained, and the trained student model is used as a trained target detection model.

2. The training method of the target detection model based on knowledge distillation according to claim 1 is characterized in that: An initial target detection model is constructed based on the YOLOv5s model, wherein the initial target detection model includes an input end, a backbone network, a neck network, and a detection end. A One-to-one Detect network structure is obtained based on the One-to-many Detect network structure of the detection end, and a One-to-one Detect network structure is added to the detection end to construct a target detection model.

3. The training method of the target detection model based on knowledge distillation according to claim 2 is characterized in that: The first loss function includes the loss function of the One-to-many Detect network structure and the loss function of the One-to-one Detect network structure; the loss function of the One-to-many Detect network structure includes the detection target confidence loss, the detection target category loss, and the prediction box loss; the loss function of the One-to-one Detect network structure includes the detection target category loss and the prediction box loss.

4. The training method of the target detection model based on knowledge distillation according to claim 1 is characterized in that: The second loss function includes a loss function of a student model and a distillation loss function; the loss function of the student model represents the error between the true label of the sample and the predicted label corresponding to the student model, and the distillation loss function represents the error between the predicted label corresponding to the teacher model and the predicted label corresponding to the student model.

5. The training method of the target detection model based on knowledge distillation according to claim 4 is characterized in that: The expression of the second loss function is as follows: Loss total2 =αLoss student +βLoss distillation , Among them, Loss total2 Represents the second loss function, Loss student Represents the loss function of the student model, Loss distillation represents the distillation loss function, and α and β are weight coefficients.

6. The training method of the target detection model based on knowledge distillation according to claim 4 is characterized in that: The loss function expression of the student model is as follows: Loss student =BCE(s,label) Among them, s is the true label of the sample, label is the predicted label corresponding to the student model, BCE(s, label) is the cross entropy loss between s and label, and z i The predicted label belonging to the i-th category output by the student model.

7. The training method of the target detection model based on knowledge distillation according to claim 5, characterized in that: The distillation loss function expression is as follows: Loss distillation =KL(q t ,q s ) Among them, q s is the predicted label output by the teacher model, q t is the predicted label output by the student model, KL(q t ,q s ) is q s With q t The KL divergence of z is is the predicted label belonging to the i-th category output by the teacher model, z it is the predicted label belonging to the i-th category output by the student model, and T is the distillation temperature.

8. The training method of the target detection model based on knowledge distillation according to claim 7, characterized in that: The parameters T, α, and β are obtained through the differential mutation algorithm: Step S301, obtain constraint conditions according to the training accuracy of the student model, and randomly generate T, α, β generating individuals that meet the constraint conditions to construct an initial population; Step S302: Calculate the loss function corresponding to each individual in the population as the fitness value of each individual, and select individuals whose fitness values ​​are less than a threshold to form an excellent population; Step S303, determine whether the termination condition is met. If so, obtain the parameters T, a, and β based on the individual with the smallest fitness value in the population; otherwise, cross and mutate the individuals in the excellent population according to the crossover probability and mutation probability to obtain new individuals, combine the individuals in the excellent population with the new individuals to obtain an updated population, use the updated population as the population for the next iteration, and repeat steps S302 to S303.

9. The method for training a target detection model based on knowledge distillation according to claim 8, characterized in that: The constraint condition expression is as follows: Among them, L is the second loss function; τ is the student model; w is the weight parameter of the student model.

10. The method for training a target detection model based on knowledge distillation according to claim 1, characterized in that: The expression of the first loss function is as follows: L total1 =L o2m +L o2o L o2m =λ obj BCE obj +λ cls BCE cls +λ box IoU L o2o =δ11BCE cls +δ2IoU, Among them, L total1 represents the first loss function, L o2m Represents the loss function of the One-to-many Detect network structure, L o2o Represents the loss function of the One-to-one Detect network structure, BCE obj To detect the target confidence loss, BCE cls is the target category loss, EIoU is the prediction box loss of the One-to-many Detect network structure, IoU is the prediction box loss of the One-to-oneDetect network structure, and λ obj , cls , box , δ1, δ2 are weight coefficients.

Images