A training method of a target detection model based on knowledge distillation

Abstract

The present application relates to a kind of training methods of target detection model based on knowledge distillation, belong to target detection field.It includes: constructing training data set and target detection model, based on training data set, target detection model is trained and outputs reasoning graph;First loss function is calculated to obtain the target detection model after training, it is used as teacher model;The parameters of teacher model are simplified to obtain student model;Second training data set is obtained based on training data set, and teacher model and student model are trained based on second training data set and export prediction label respectively;Second loss function is calculated to obtain the student model after training as the target detection model after training.The method adopts difference variation algorithm to obtain the optimal distillation temperature and loss coefficient, and the student model is trained by knowledge distillation to obtain the target detection model after training, solve the problem that the existing target detection model is in the case of small amount of sample training, leading to low detection precision.

Description

TECHNICAL FIELD

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

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

[0003] However, the existing target detection model based on deep learning usually needs a large number of samples for training, and in the case of a small number of samples, the detection accuracy is usually low. Moreover, the existing target detection model based on deep learning is usually not simplified, which is not conducive to deploying the model under the condition of limited hardware resources.

[0004] Therefore, it is an urgent problem to provide a method for training a target detection model in the case of a small number of samples. SUMMARY

[0005] In view of the above analysis, the present application aims to provide a training method of a target detection model based on knowledge distillation, in order to solve the problem of low detection accuracy of the existing target detection model in the case of a small number of samples.

[0006] The present application provides a training method of a target detection model based on knowledge distillation, which comprises the following steps:

[0007] The images in the public data set are preprocessed and labeled to construct a training data set; a target detection model is constructed, each sample in the training data set is input into the target detection model for training, and the corresponding inference graph is output; the inference graph includes a prediction box;

[0008] A first loss function is calculated by the error between the labeled box of the sample and the prediction box of the corresponding inference graph, and when the first loss function reaches a threshold value, a trained target detection model is obtained, which is used as a teacher model; the convolution channel number of the backbone network of the teacher model is halved to obtain a student model;

[0009] Based on the training data set, a second training data set is obtained, 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 output corresponding prediction labels respectively; a second loss function is calculated through the error between the real label of the sample and the prediction label corresponding to the student model and the error between the prediction label corresponding to the teacher model and the prediction label corresponding to the student model, and when the second loss function reaches a threshold value, a trained student model is obtained, and the trained student model is used as a trained target detection model.

[0010] Further, an initial target detection model is constructed based on the YOLOv5s model, the initial target detection model comprising 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 the One-to-one Detect network structure is added to the detection end, so as to construct the target detection model.

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

[0012] Further, the second loss function comprises a loss function of the student model and a distillation loss function; the loss function of the student model represents the error between the real label of the sample and the prediction label corresponding to the student model, and the distillation loss function represents the error between the prediction label corresponding to the teacher model and the prediction label corresponding to the student model.

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

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

[0015] Wherein, 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] Further, the expression of the loss function of the student model is as follows:

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

[0018]

[0019] where s is the real label of the sample, label is the predicted label corresponding to the student model, BCE(s, label) is the cross-entropy loss of s and label, z i is the predicted label of the i-th class output by the student model.

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

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

[0022]

[0023] where 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 the KL divergence of q s and q t , z is is the predicted label of the i-th class output by the teacher model, z it is the predicted label of the i-th class output by the student model, and T is the distillation temperature.

[0024] Further, the parameters T, a and b are obtained by a differential mutation algorithm:

[0025] In step S301, the constraint condition is obtained according to the training accuracy of the student model, T, a and b satisfying the constraint condition are randomly generated to generate individuals to construct an initial population;

[0026] In step S302, the loss function corresponding to each individual in the population is calculated as the fitness value of each individual, and the individuals with fitness values less than a threshold value are selected to form an excellent population.

[0027] In step S303, it is judged whether the termination condition is reached, if yes, the parameters T, a and b are obtained according to the individual with the minimum fitness value in the population; otherwise, the individuals in the excellent population are crossed and mutated according to the crossover probability and mutation probability to obtain new individuals, the individuals in the excellent population and the new individuals are combined to obtain an updated population, the updated population is taken as the population of the next iteration, and steps S302 to S303 are repeatedly executed.

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

[0029]

[0030] Wherein, L is a second loss function; Tau is a student model; W is a weight parameter of the student model.

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

[0032] L total1 = L o2m + L o2o

[0033] L o2m = lambda obj BCE obj + lambda cls BCE cls + lambda box EIoU

[0034] L o2o = delta1 cls BCE + delta2

[0035] Wherein, 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 is the detection target confidence loss, BCE cls is the detection target class 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, lambda obj , lambda cls , lambda box , delta1 and delta2 are weight coefficients.

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

[0037] 1. The trained target detection model is obtained by training the student model through knowledge distillation, which can restore the accuracy of the lightweight model under the condition of using only a small amount of samples, and is convenient for deploying the model under the condition of limited hardware resources; The problem of low detection precision of the existing target detection model under the condition of training with a small amount of samples is solved.

[0038] 2、The differential variation algorithm is adopted to obtain optimal distillation temperature, student model loss coefficient alpha and distillation loss coefficient beta, and the prediction accuracy of the student model is improved.

[0039] 3、The detection end is improved, through the One-to-one Detect network structure, the prediction result can be directly outputted during use, without post-processing of the inference graph, and the end-to-end target detection effect is generated.

[0040] 4、The C3AM module adds an attention mechanism, prunes the convolution channel to be fused through channel information enhancement, removes the redundant channel in the network, reduces the parameters of the model, and improves the efficiency of model training and inference.

[0041] In the present application, the above technical solutions can be combined with each other to realize more preferred combination solutions. Other features and advantages of the present application will be described in the subsequent specification, and some advantages will become apparent from the specification or by implementing the present application. The purpose and other advantages of the present application can be realized and obtained through the contents specifically pointed out in the specification and the drawings. BRIEF DESCRIPTION OF DRAWINGS

[0042] The accompanying drawings are included to provide a further understanding of the embodiments, and are not considered limiting of the application. In the drawings, like reference numerals indicate like parts throughout the several views.

[0043] Figure 1 The flow chart of the training method of the target detection model based on knowledge distillation of the embodiment of the present application is shown in Figure 1.

[0044] Figure 2(a) is a schematic diagram of the overall network structure of the initial target detection model of the embodiment of the present application.

[0045] Figure 2(b) is a schematic diagram of the network structure of the CBS, C3_1, C3_2, BottleNeck1 and SPPF modules of the initial target detection model of the embodiment of the present application.

[0046] Figure 3 Figure 3 is a schematic diagram of the structure of the C3AM module of the embodiment of the present application.

[0047] Figure 4 Figure 4 is a schematic diagram of the structure of the improved neck network of the embodiment of the present application.

[0048] Figure 5 Figure 5 is a schematic diagram of the network structure of the improved target detection model of the embodiment of the present application. DETAILED DESCRIPTION

[0049] Preferred embodiments of the present application will be described in detail below with reference to the drawings, wherein the drawings constitute a part of this application, and are used to explain the principles of the application together with the embodiments of the application, but are not used to limit the scope of the application.

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

[0051] Step S1, pre-process the images in the public data set, and construct a training data set after labeling; construct a target detection model, input each sample in the training data set into the target detection model for training, and output the corresponding inference graph; the inference graph comprises a prediction box;

[0052] Step S2, calculate a first loss function through the error between the labeled box of the sample and the prediction box of the corresponding inference graph, obtain the trained target detection model when the first loss function reaches a threshold value, and take it as a teacher model; halve the convolution channel number of the backbone network of the teacher model to obtain a student model;

[0053] Step S3, obtain a second training data set based on the training data set, input each sample in the second training data set into the teacher model and the student model for training respectively, and the teacher model and the student model output corresponding prediction labels respectively; calculate a second loss function through the error between the real label of the sample and the prediction label of the student model, and the error between the prediction label of the teacher model and the prediction label of the student model; obtain a trained student model when the second loss function reaches a threshold value, and take the trained student model as a trained target detection model.

[0054] Specifically, in step S1, the pre-processing of the images in the public data set and the construction of the training data set after labeling comprises:

[0055] Obtain a public data set, and randomly select part of the images therefrom;

[0056] Pre-process the selected images by removing the mean, converting the image channel, and converting the image format;

[0057] Adaptively scale the pre-processed images to a uniform size;

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

[0059] Label the adaptively scaled and image-enhanced images to construct a training data set.

[0060] ​Specifically, a public data set, such as a COCO data set, a visdrone data set, and a vehicules data set, is acquired, and part of the images are randomly selected therefrom. The de-meaning refers to subtracting the mean values of the corresponding channels of the ImageNet data set from the pixel values of the R, G, and B channels of the selected images, respectively. The converting image channel refers to uniformly converting the channel format of the selected images into an RGB format. The converting image format refers to converting the format of the selected images into the same format, such as a jpeg format.

[0061] It can be understood that the de-meaning of the selected images can reduce the deviation of the pixel distribution of the input images.

[0062] Further, the self-adaptive scaling of the preprocessed images to a uniform size comprises:

[0063] setting a target size; for each preprocessed image:

[0064] selecting the size of the larger one from the width and the height of the image, taking the ratio of the target size to the selected width or height of the image 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 scaled image is different from the target size, filling the four sides of the scaled image with a fixed pixel value according to the difference between the size of the scaled image and the target size, so that the size of the filled image is consistent with the target size and the central content remains unchanged.

[0067] For example, the target size is set to 640x640, and assuming that the size of an image is 1280x720, the width 1280 of the image is selected, and the ratio of the width 640 of the target size to 1280 is taken as the scaling ratio of the image. For example, the target size is set to 640x640, and assuming that the size of an image is 1280x720, the width 1280 of the image is selected, and the ratio of the width 640 of the target size to 1280 is taken as the scaling ratio of the image. For example, the target size is set to 640x640, and assuming that the size of an image is 1280x720, the width 1280 of the image is selected, and the ratio of the width 640 of the target size to 1280 is taken as the scaling ratio of the image.

[0068] It should be noted that the preprocessed images are scaled, and the four sides of the scaled image are filled with a fixed pixel value, so that the preprocessed images are self-adaptively scaled to a uniform size and the content of the original image is centered.

[0069] Further, the image enhancement processing of the self-adaptively scaled images comprises one or more of the following methods:

[0070] randomly selecting part of the content in the image and copying to other images to generate new images;

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

[0072] randomly selecting two images in the image, and mixing the two images according to a certain proportion to generate a new image;

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

[0074] randomly occluding part of the area in the image to generate a new image.

[0075] It should be noted that skewing includes horizontal skewing and vertical skewing. Horizontal skewing achieves a distortion effect by adding a horizontal displacement amount on each row of the image, and vertical skewing achieves a distortion effect by adding a vertical displacement amount 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, thereby providing a basis for improving the robustness and generalization ability of the model.

[0077] Further, the LabelImg labeling tool is used to label each target to be detected in the image after adaptive scaling and image enhancement with a rectangular box to construct a data set; the target to be detected is selected according to needs; the labeling content includes a class number, coordinates of a center point of the target to be detected, a height, and a width of the labeling box.

[0078] Specifically, the training data set includes images after adaptive scaling and image enhancement and corresponding labeling content thereof.

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

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

[0081] Specifically, the initial target detection model is constructed based on a YOLOv5s model, and a network structure diagram of the initial target detection model is as shown in FIG. 2(a) and FIG. 2(b). The input end is configured to obtain an image to be detected. The backbone network is configured to extract feature information of different scales from the image to obtain feature maps of corresponding scales. The neck network includes a feature pyramid network (FPN) and a path aggregation network (PAN), which are configured to fuse the feature maps of corresponding scales, fuse shallow feature information and deep feature information of the image, and input to the detection end. The detection end is configured to output target detection results of corresponding scales. The CBS module includes a convolution layer (Conv2d), a batch normalization layer (BN), and a SiLu activation function connected in sequence, which are configured to extract feature information of different scales from 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, which are configured to fuse feature information of different scales. 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 skip connection and a double CBS module connected in parallel. The double CBS module includes a fourth CBS module 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, which are configured to fuse feature information of 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. The second bottleneck layer includes a fourth CBS module 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, which are configured to extract global feature information of the image. The output of the first CBS module and the output of the three maximum pooling layers are further connected to the feature fusion layer in a skip connection manner.

[0082] Further, as Figure 3As shown, the C3AM module includes a Shuffle module, a split layer (Split), first to fourth CBS modules, a feature fusion layer, and an AM module. The input feature map is divided into two paths by the Shuffle module and the split layer connected in sequence, and the number of channels of each path is reduced to half of the original number of channels. One path is input to the feature fusion layer after convolution operation by the first CBS module, and the other path is input to the feature fusion layer after convolution operation by the second CBS module. The output of the second CBS module is also input to the feature fusion layer after convolution operation by the third CBS module, and the output of the third CBS module is also input to the feature fusion layer after convolution operation by the fourth CBS module. 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 used to obtain a feature map with enhanced channel information.

[0084] Further, the input of the fused feature map to the AM module for attention calculation to obtain a channel attention weighted feature map includes:

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

[0086] Step S102, calculate the maximum feature value and the average feature value of the feature map of each channel after channel removal, generate a global maximum feature vector based on the maximum feature value, and generate an average feature vector based on the average feature value;

[0087] Step S103, perform convolution operation on the global maximum feature vector and the average feature vector with a 1x1 convolution kernel respectively, to learn the attention weight of each channel after channel removal; add the global maximum feature vector and the average feature vector after convolution operation to obtain a sum vector;

[0088] Step S104, generate a channel attention weight vector based on the sum vector by a Sigmoid activation function;

[0089] Step S105, the channel attention weight values in the channel attention weight vector are sorted from high to low, a pruning ratio is set, the lowest part of the channel attention weight values is set to 0 according to the pruning ratio 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 corresponding to the channel to obtain a channel attention weighted feature map.

[0090] For example, assuming that the i-th channel fused feature map is F i ∈R C×H×W , wherein 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; pHash calculation is performed on each channel fused feature map to obtain a corresponding hash value; if the Hamming distance between the hash values corresponding to two channel fused feature maps is less than 5, one of the two channels is randomly retained and the other channel is removed to obtain a channel removed i-th channel fused feature map F' i ∈R C′×H×W , wherein C' is the number of channels after channel removal, and C' ≤ C; the maximum eigenvalue and the average eigenvalue of the feature map of each channel after channel removal are calculated, a global maximum eigenvalue vector e max ∈R C′×1×1 is generated based on the average eigenvalue of the feature map of each channel. min ∈R C′×1×1 , wherein e max = MaxPool(F'), and e min = AvgPool(F'); the global maximum eigenvalue vector e max and the average eigenvalue vector e min are respectively subjected to convolution operation with a 1x1 convolution kernel to obtain e max ' = Conv 1×1 (e max ), and e min ' = Conv 1×1 (e min ), e max ' and e min ' are added to obtain a sum vector V = e max ' + e min ', V ∈ R C′×1×1 ; a channel attention weight vector The channel attention weight values in the channel attention weight vector are sorted from high to low, a 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 the corresponding channel to obtain a channel attention weighted feature map.

[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 value, it is considered that the feature maps of the two channels are similar, and therefore one of the two channels is randomly retained and the other channel is removed to reduce the number of channels. The global maximum feature vector contains the maximum feature values of the feature maps of each channel after channel removal, and the average feature vector contains the average feature values of the feature maps of each channel after channel removal. When a channel attention weight value is low, it is set to 0 according to the pruning ratio, and at this time, the multiplication of the channel attention weight value and the feature map of the corresponding channel is also 0, so that the corresponding channel is inhibited, thereby enhancing the channel information with high correlation and inhibiting the channel information with low correlation.

[0092] Further, the One-to-many Detect network structure based on the detection end comprises:

[0093] Based on the training data set, the bounding boxes are clustered to obtain a plurality of anchor boxes corresponding to large, medium and small targets respectively;

[0094] The One-to-many Detect network structure of the detection end comprises a plurality of convolutional layers for predicting a plurality of prediction boxes by using a plurality of anchor boxes corresponding to large, medium and small targets respectively;

[0095] Based on the One-to-many Detect network structure of the detection end, a One-to-one Detect network structure with the same structure is constructed; the One-to-one Detect network structure is used to select the anchor box closest to the corresponding target from the plurality of anchor boxes corresponding to each target respectively for prediction, and each target obtains one prediction box.

[0096] Further, the clustering of the bounding boxes of the training data set to obtain a plurality of anchor boxes corresponding to large, medium and small targets respectively comprises: obtaining the size of all the bounding boxes of the training data set, i.e. the width and height of the bounding boxes, using a clustering algorithm to analyze the size of all the bounding boxes, and obtaining anchor boxes corresponding to large, medium and small targets respectively; wherein each target corresponds to a plurality of anchor boxes with different sizes.

[0097] Exemplarily, the sizes of all the labeled boxes are clustered using a K-means clustering algorithm, and three different sizes of anchor boxes corresponding to large, medium and small targets are obtained respectively, so that nine different sizes of anchor boxes corresponding to large, medium and small targets are obtained.

[0098] Further, the One-to-many Detect network structure of the detection end is used to predict each target using a plurality of anchor boxes corresponding to the target respectively, and each target obtains a plurality of predicted boxes.

[0099] Exemplarily, the One-to-many Detect network structure of the detection end predicts on feature maps with sizes of 80x80, 40x40 and 20x20 respectively, for example, three anchor boxes corresponding to large targets are used to predict on the feature map with the size of 80x80, and each large target obtains three corresponding predicted boxes.

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

[0101] It can be understood that when the One-to-many Detect network structure is used for prediction, a large number of target predicted boxes are obtained on the inference graph, and many of the predicted boxes are redundant, so it is necessary to remove the redundant predicted boxes by using the NMS (Non-Maximum Suppression) method, that is, the inference graph needs to be post-processed to obtain the predicted target.

[0102] Further, the selecting of the anchor box closest to the corresponding target from the plurality of anchor boxes corresponding to the large, medium and small targets respectively includes: calculating the aspect ratio of each anchor box corresponding to the corresponding target, calculating the aspect ratio of the labeled box of the corresponding target, and selecting the anchor box closest to the aspect ratio of the corresponding target from the plurality of anchor boxes.

[0103] Further, the One-to-one Detect network structure obtains the predicted box of each target by 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 Let b be the x-coordinate of the center point of the prediction box. y Let b be the y-coordinate of the center point of the prediction box. w b is the width of the prediction box h t is the height of the predicted bounding box. x t represents the x-coordinate of the center point of the predicted bounding box output by the model. y t represents the y-coordinate of the center point of the predicted bounding box output by the model. w t represents the width of the predicted bounding box output by the model. h c represents the height of the predicted bounding box output by the model. x c is the x-coordinate of the top-left corner of the grid cell containing the center point of the target annotation box. y p is the y-coordinate of the top-left corner of the grid cell containing the center point of the target annotation box. w p is the width of the closest selected anchor frame. h The height of the closest selected anchor frame.

[0109] Furthermore, such as 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 at the detection end to construct the target detection model. A schematic diagram of the network structure of the target detection model is shown below. Figure 5 As shown.

[0110] Understandably, the number of convolutional channels in the initial object detection model is usually determined manually based on experience, resulting in redundancy. This significantly impacts the training and inference efficiency of the model, especially under conditions of limited hardware resources. By employing the C3AM module, the convolutional channels to be fused are pruned, reducing model parameters and improving the efficiency of model training and inference. This invention, through a One-to-one Detect network structure, can directly output prediction results without post-processing the inference graph, producing end-to-end object detection results. This improves inference efficiency while maintaining the model's prediction accuracy; it also solves the problem of low detection speed caused by post-processing of detection results in existing object detection methods.

[0111] Further, the inputting each sample in the training data set into the target detection model for training and outputting corresponding inference graphs comprises: the One-to-many Detect network structure and the One-to-one Detect network structure output corresponding inference graphs respectively.

[0112] Specifically, in step S2, the first loss function comprises a loss function of the One-to-many Detect network structure and a loss function of the One-to-one Detect network structure; the loss function of the One-to-many Detect network structure comprises 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 comprises a detection target category loss and a prediction box loss; and an 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 EIoU

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

[0116] wherein, 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 is the detection target confidence loss, BCE cls is the detection 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, λ obj , λ cls , λ box , δ1 and δ2 are weight coefficients.

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

[0118] Further, the detection target confidence loss represents the error of prediction accuracy, i.e. whether there is a target in the prediction box, and the loss function expression of the detection target confidence loss is as follows:

[0119]

[0120] Wherein, N is the total number of samples, m is the number of positive samples, N-m is the number of negative samples, y i =1 indicates that there is a target in the prediction box, y i =0 indicates that there is no target in the prediction box, p1 is the probability of the presence of the target in the prediction box, and a is a coefficient.

[0121] It can be understood that the negative sample corresponds to the image background, and the positive sample corresponds to the actual target region in the image. Generally, the number of negative samples is more, and the number of positive samples is less, which can easily lead to the loss value of the negative sample to submerge the loss value of the positive sample, and reduce the convergence speed of the model and the accuracy of target detection. The BCE obj loss function proposed in the present application can dynamically adjust the loss value of the positive and negative samples, thereby improving the convergence speed of the model and the accuracy of target detection.

[0122] Further, the detection target category loss represents the error between the category of the prediction box and the category of the labeled box, and the loss function expression of the detection target category loss is as follows:

[0123]

[0124] Wherein, N p is the number of positive samples, n is the number of categories, C i =1 indicates that the predicted category belongs to C i category, C i =0 indicates that the predicted category does not belong to C i category, and p i represents the probability of predicting the C i category.

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

[0126]

[0127] Wherein, b pre represents the area occupied by the prediction box, b gt represents the area occupied by the labeled box, IoU is the intersection over union of the prediction box and the labeled box, and p 2 (bpre , b gt represents the distance between the center point of the prediction box and the center point of the labeled box, w pre is the width of the prediction box, w gt is the width of the labeled box, p 2 (w pre , w gt ) represents the square of the difference between the width of the prediction box and the width of the labeled box, h pre is the height of the prediction box, h gt is the height of the labeled box, p 2 (h pre , h gt ) represents the square of the difference between the height of the prediction box and the height of the labeled box, c w 2 represents the square of the minimum circumscribed rectangle width of the prediction box and the labeled box, c h 2 represents the square of the minimum circumscribed rectangle height of the prediction box and the labeled box.

[0128] It should be noted that the positive sample refers to the error between the prediction box in the inference graph and the actual labeled box of the sample being less than a threshold value, indicating that the area intersection over union of the prediction box and the labeled box is greater than a threshold value and the class of the prediction box is consistent with the class 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 of the prediction box and the labeled box.

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

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

[0131] Further, 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 corresponding predicted label of the student model, and the distillation loss function represents the error between the corresponding predicted label of the teacher model and the corresponding predicted label of the student model; the expression of the second loss function is as follows:

[0132] Loss total2 = aLoss student + bLosd distillation ,

[0133] wherein 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 a and β are weight coefficients.

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

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

[0136]

[0137] wherein s is the real label of the sample, label is the predicted label corresponding to the student model, BCS(s, label) is the cross-entropy loss of s and label, z i is the predicted label of the student model belonging to the i-th class.

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

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

[0140]

[0141]

[0142] wherein 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 the KL divergence of q s and q t , z is is the predicted label of the teacher model belonging to the i-th class, z it is the predicted label of the student model belonging to the i-th class, and T is the distillation temperature.

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

[0144] In step S301, the constraint condition is obtained according to the training accuracy of the student model, T, a and β satisfying the constraint condition are randomly generated to generate individuals and construct an initial population;

[0145] In step S302, the loss function corresponding to each individual in the population is calculated as the fitness value of each individual, and the individuals with fitness values less than a threshold value are selected to form an excellent population.

[0146] Step S303, it is judged whether the termination condition is reached, if yes, the parameters T, a, b are obtained according to the individual with the minimum fitness value in the population; otherwise, the individuals in the excellent population are crossed and mutated according to the crossover probability and the mutation probability to obtain new individuals, the individuals in the excellent population and the new individuals are combined to obtain an updated population, and the updated population is taken as the population of the next iteration, and steps S302 to S303 are repeatedly executed.

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

[0148]

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

[0150] Specifically, T, a, b satisfying the constraint condition are randomly generated by the following formula:

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

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

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

[0154] Wherein, rand(0, 1) represents a random decimal number between 0 and 1, i represents the serial number of the individual in the first population, and j represents the length of the individual chromosome.

[0155] Specifically, each individual is represented as X = [T, a, b], and the individual is generated by randomly generating a value satisfying the constraint condition in the value range of T, a, b according to the above formula to complete the construction of the initial population.

[0156] Specifically, in step S302, the loss function Loss total2 of each individual in the population is calculated 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 minimum fitness value in the current population is less than a threshold value or the number of iterations reaches a maximum value.

[0158] Further, if the termination condition is not reached, the individuals in the excellent population are randomly selected according to the mutation probability to generate new individuals by the following formula:

[0159]

[0160] Wherein, F represents the mutation probability; p1, p2, p3 are random integers, representing the serial number of the individual in the excellent population; t represents the current iteration number, h ij (t+1) represents the newly generated individual;

[0161] Specifically, if the termination condition is not reached, the individual is randomly selected in the excellent population according to the crossover probability to generate a new individual through the following formula:

[0162]

[0163] Wherein, randl ij is a random number between 0 and 1, CR is the crossover probability, x ij (t) is the individual randomly selected in the excellent population.

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

[0165] Further, the image to be detected is input into the trained target detection model, the One-to-many Detect network structure is frozen, the prediction label is directly output through the One-to-one Detect network structure, and it is judged whether the prediction label is the target to be detected, so as to realize target recognition.

[0166] It can be understood that the selection of the three key parameters of distillation temperature, distillation loss coefficient and student model loss coefficient is usually determined by artificial experience, which seriously affects the learning effect of the actual student model. The differential mutation algorithm is used to obtain the optimal distillation temperature, student model loss coefficient alpha and distillation loss coefficient beta, which improves the prediction accuracy of the student model. The trained target detection model is obtained by training the student model through knowledge distillation. In the case of using only a small amount of samples, the accuracy of the lightweight model can be restored, which is convenient for deploying the model under the condition of limited hardware resources; solve the problem that the existing target detection model leads to low detection accuracy in the case of training with a small amount of samples.

[0167] Performance evaluation:

[0168] Specifically, the performance evaluation index is set to evaluate the performance of the improved target detection model, and the performance evaluation index includes mean average precision (mAP), precision (P) and recall (R). The precision P is used to evaluate the accuracy of target detection, the recall R is used to evaluate the comprehensiveness of target detection, and the mAP is the average value of the area under the PR curve of all classes, which comprehensively considers the precision and recall of the model.

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

[0170]

[0171] Wherein, TP represents the number of correct target detection results, FP represents the number of incorrect target detection results, FN represents the number of missed detection of the target, AP represents the area under the PR curve of the category, n c is the total number of categories.

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

[0173] Table 1 Experimental results

[0174]

[0175] As shown in the above table, compared with YOLOv5s, the improved target detection model through the method proposed in the application greatly reduces the parameter amount while also improving the average precision to some 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 application has the following beneficial effects:

[0177] 1. The trained target detection model is obtained by training the student model through knowledge distillation, which can restore the accuracy of the lightweight model under the condition of using only a small amount of samples, and is convenient for deploying the model under the condition of limited hardware resources; solve the problem of low detection accuracy of the existing target detection model under the condition of a small amount of sample training.

[0178] 2. The differential mutation algorithm is used to obtain the optimal distillation temperature, student model loss coefficient a and distillation loss coefficient β, which improves the prediction accuracy of the student model.

[0179] 3. The detection end is improved in the application, and through the One-to-one Detect network structure, the prediction result can be directly output during use, without post-processing of the inference graph, and an end-to-end target detection effect is generated.

[0180] 4. The C3AM module proposed in the application adds an attention mechanism, prunes the convolution 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 inference.

[0181] Those skilled in the art can understand that all or part of the processes of the above-mentioned embodiment methods can be completed by instructing the relevant hardware through a computer program, and the program can be stored in a computer readable storage medium. The computer readable storage medium is a disk, an optical disk, a read-only memory, a random access memory, etc.

[0182] The above description is only the preferred embodiment of the present application, but the protection scope of the present application is not limited to this. Any changes or replacements within the technical scope disclosed by the present application can be easily thought by those skilled in the art, and should be covered within the protection scope of the present application.

Claims

1. A method for training a target detection model based on knowledge distillation, characterized in that, The method comprises the following steps: Preprocess and label images in a public data set to construct a training data set; construct a target detection model, input each sample in the training data set into the target detection model for training, and output a corresponding inference graph; the inference graph comprises a prediction box; An initial target detection model is constructed based on a YOLOv5s model, and the initial target detection model comprises an input end, a backbone network, a neck network and a detection end; A first loss function is calculated based on the error between the label box of a sample and the prediction box of the corresponding inference graph, and when the first loss function reaches a threshold, a trained target detection model is obtained, which is used as a teacher model; the convolution channel number of the backbone network of the teacher model is halved to obtain a student model; 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 EIoU L o2o = δ11BCE cls + δ2IoU, wherein, 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 is the detection target confidence loss, BCE cls is the detection target class 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, λ obj , λ cls , λ box , δ1, δ2 are weight coefficients; A second training data set is obtained based on the training data set, 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 output corresponding prediction labels respectively; a second loss function is calculated based on the error between the real label of a sample and the prediction label of the student model and the error between the prediction label of the teacher model and the prediction label of the student model, and when the second loss function reaches a threshold, a trained student model is obtained, which is used as a trained target detection model.

2. The training method of a knowledge distillation based target detection model according to claim 1, characterized in that, The second loss function comprises a loss function of the student model and a distillation loss function; the loss function of the student model represents the error between the real label of a sample and the prediction label of the student model, and the distillation loss function represents the error between the prediction label of the teacher model and the prediction label of the student model.

3. The training method of a knowledge distillation based target detection model according to claim 2, characterized in that, The expression of the second loss function is as follows: Loss total2 = aLoss studen t + bLoss distillartion , wherein Loss total2 represents a second loss function, Loss student represents a loss function of the student model, Loss distillation represents a distillation loss function, and a and β are weight coefficients.

4. The training method of a knowledge distillation based target detection model according to claim 2, characterized in that, The expression of the loss function of the student model is as follows: Loss student = BCE(s, label) where 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 of s and label, z i is the predicted label of the i-th class output by the student model. 5.The method of claim 3, wherein, The expression of the distillation loss function is as follows: Loss distillation = KL(q t , q s ) where 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 the KL divergence of q s and q t , z is is the predicted label output by the teacher model belonging to the i-th class, z it is the predicted label output by the student model belonging to the i-th class, and T is the distillation temperature. 6.The method of claim 5, wherein, Parameters T, α and β are obtained through a differential mutation algorithm: In step S301, a constraint condition is obtained according to the training accuracy of the student model, T, α and β that meet the constraint condition are randomly generated to generate individuals and construct an initial population; In step S302, the loss function corresponding to each individual in the population is calculated as the fitness value of each individual, and individuals with fitness values less than a threshold are selected to form an excellent population; In step S303, it is determined whether a termination condition is reached, if yes, parameters T, α and β are obtained according to the individual with the minimum fitness value in the population; otherwise, individuals in the excellent population are crossed and mutated according to a crossover probability and a mutation probability to obtain new individuals, the individuals in the excellent population and the new individuals are combined to obtain an updated population, and the updated population is used as the population of the next iteration, and steps S302 to S303 are repeatedly executed.

7. The training method of a knowledge distillation based target detection model according to claim 6, characterized in that, The expression of the constraint condition is as follows: Wherein, L is the second loss function; τ is the student model; and w is the weight parameter of the student model.

Images