Model training method and device based on knowledge distillation, and electronic equipment
By aligning the feature layers and output layers of the teacher model and the student model, determining the feature distillation loss and output distillation loss, and improving the loss function of the student model, the problem of mismatch between prediction box size, class score, and feature size in heterogeneous model training is solved, enabling efficient training of the student model, which has similar capabilities and broad application scenarios as the teacher model.
Patent Information
- Application Number
- CN202211608782.8
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2022-12-14
- Publication Date
- 2025-12-09
- Estimated Expiration
- 2042-12-14
AI Technical Summary
In existing technologies, there are no complete solutions to problems such as mismatch in prediction box size, class score meaning, and feature scale for heterogeneous models. Furthermore, existing technologies do not fully address the issues of mismatch in prediction box size, class score meaning, and feature scale during the training process of heterogeneous models.
By aligning the feature and output layers of the teacher model with those of the student model, determining the feature distillation loss and output distillation loss, improving the loss function of the student model, guiding the training of the student model, and achieving information alignment of heterogeneous models in the knowledge distillation process.
Without altering the student model structure, it can be made to have generalization capabilities and functions similar to the teacher model, reducing deployment resource requirements and expanding application scenarios.
Smart Images

Figure CN115953643B_ABST
Abstract
Description
TECHNICAL FIELD
[0001] The present application relates to the technical field of machine learning, and in particular to a model training method and device based on knowledge distillation and an electronic device. BACKGROUND
[0002] The task of object detection is to find the objects of interest in an image and determine their locations and categories. The way to achieve object detection is usually to train an object detection model and achieve the object detection function through the trained object detection model. For the trained object detection model, the complexity of the model usually increases with the improvement of the accuracy of the model, and thus the requirement for deployment resources will be higher. Knowledge distillation is a model compression method that induces student model training by introducing a teacher model, which can achieve knowledge transfer from the teacher model to the student model. Therefore, through the method of knowledge distillation, the complex model trained can be used as a teacher model, and the knowledge of the teacher model can be transferred to a student model with simpler structure, so that the student model can retain the accuracy of the teacher model while solving the problem of insufficient deployment resources.
[0003] The teacher model and the student model mentioned above based on knowledge distillation can be isomorphic models or heterogeneous models. In the present application, the teacher model of the two-stage detection network and the student model of the single-stage detection network used in the process of training the object detection model are heterogeneous models. The problems of mismatched prediction box size, mismatched category score significance, and mismatched feature size caused by the heterogeneous model have not yet been completely solved. Therefore, the present application proposes a model training method and device based on knowledge distillation, an electronic device, and a storage medium to solve the above-mentioned series of heterogeneous distillation problems. SUMMARY
[0004] The present application provides a model training method and device based on knowledge distillation and an electronic device, which can solve the series of heterogeneous distillation problems such as mismatched prediction box size, mismatched category score significance, and mismatched feature size when training an object detection model using the method of knowledge distillation, where the student model of the single-stage detection network and the teacher model of the two-stage detection network are heterogeneous models.
[0005] To achieve the above-mentioned purpose, the embodiments of the present application adopt the following technical solutions:
[0006] In a first aspect, the application provides a model training method based on knowledge distillation. In the method, a teacher model involved in knowledge distillation is a two-stage detection network, and a student model is a single-stage detection network, i.e., the teacher model and the student model are heterogeneous models. Knowledge distillation in the teacher model is migrated to the student model. The student model after the knowledge distillation process is a trained model, which is used for target detection on an input image. The method comprises the following steps: obtaining a trained teacher model; aligning a feature scale of a feature layer of the student model with a feature scale of a feature layer of the teacher model, and determining a feature distillation loss. The feature distillation loss represents a difference degree of the teacher model and the student model at the feature layer after the alignment of the feature scales. Aligning a prediction box and a corresponding probability distribution of an output layer of the student model with a prediction box and a corresponding probability distribution of an output layer of the teacher model, and determining an output distillation loss. The output distillation loss represents a difference degree of the teacher model and the student model at the output layer after the alignment of the prediction box and the corresponding probability distribution. The probability distribution is a probability distribution of an output of the prediction box. Improving a loss function of the student model according to the feature distillation loss and the output distillation loss. Training the student model based on the improved loss function to obtain a trained model.
[0007] The technical scheme provided by the embodiments of the application brings at least the following beneficial effects:
[0008] The student model for target detection is trained in a manner of knowledge distillation, and the teacher model involved in the knowledge distillation is a two-stage detection network, and the student model is a single-stage detection network, that is, the teacher model and the student model are heterogeneous models. When training, the application aligns the prediction frame and the corresponding probability distribution of the output layer of the student model with the prediction frame and the corresponding probability distribution of the output layer of the teacher model, thereby solving the problem of misalignment of the output information of the student model and the teacher model. The feature scale of the feature layer of the student model is aligned with the feature scale of the feature layer of the teacher model, thereby solving the problem of misalignment of the feature information of the student model and the teacher model. The loss function of the student model is improved by determining the feature distillation loss and the output distillation loss, and the training of the student model is guided by the improved loss function, so that the student model learns the knowledge of the teacher model, and the student model has the generalization ability and function close to the teacher model without changing the structure of the student model. In summary, the application proposes a complete heterogeneous distillation scheme. In the case that the teacher model is a two-stage detection network and the student model is a single-stage detection network, that is, the teacher model and the student model are heterogeneous models, the training of the target detection model is realized in a manner of knowledge distillation. On the other hand, the student model is distilled by the teacher model with a more complex structure and a more perfect function, so that the student model has a similar ability to the teacher model without changing the structure of the student model. At the same time, the student model has a simple structure and low requirements for deployment resource performance, so that the application scenarios of the student model for target detection obtained by training are more extensive.
[0009] In a possible implementation, the alignment of the feature scales of the student feature set and the teacher feature set includes: obtaining a teacher feature set of the teacher model; wherein the teacher feature in the teacher feature set is a feature of a target region in a first image input to the teacher model, and the target region is a region determined based on an anchor box, and the feature of the target region includes a foreground region feature and a background region feature of the target region; obtaining a student feature set in the student model; wherein the student feature in the student feature set is a feature of the target region in the student model; and converting the feature scales in the teacher feature set and the student feature set to the same dimension.
[0010] In this possible implementation, a method for solving the problem of mismatching of the feature scales of the teacher model and the student model in the knowledge distillation process is proposed, and the feature information of the student model and the teacher model is converted to the same dimension, which helps to improve the implementability of the scheme.
[0011] In a possible implementation, the student feature set obtained in the features of each layer of the student model includes: extracting student region features of each layer based on a target region in the first image of the input student model to obtain the student feature set.
[0012] In a possible implementation, the feature distillation loss at the feature layer is determined by: calculating feature similarities between each teacher feature in the teacher feature set and each student feature in the student feature set based on an attention mechanism or a cosine similarity algorithm, and determining the feature distillation loss according to the feature similarities.
[0013] In this possible implementation, the feature similarities between each feature in the teacher feature set and each feature in the student feature set can be obtained, and based on this, the feature distillation loss in the feature alignment process can be further calculated, so that the loss function of the student model is improved based on the feature distillation loss, and the student model obtained by training has better target detection performance.
[0014] In a possible implementation, the alignment of the prediction box and the corresponding probability distribution of the output layer of the student model with the prediction box and the corresponding probability distribution of the output layer of the teacher model includes: calculating an intersection over union of the prediction box of the output layer of the student model and the prediction box of the output layer of the teacher model, and selecting a target prediction box that matches the prediction box of the output layer of the teacher model from the prediction boxes of the output layer of the student model according to the intersection over union; obtaining a target probability distribution corresponding to the target prediction box, the target probability distribution being the probability distribution output by the target prediction box; and aligning the target probability distribution with the probability distribution corresponding to the prediction box of the output layer of the teacher model.
[0015] In this possible implementation, by calculating the intersection over union of each teacher prediction box and student prediction box, the student prediction box that best matches the teacher prediction box is selected, and the probability distribution of the selected student prediction box is aligned with the probability distribution of the corresponding teacher prediction box, thereby solving the problem of inconsistent prediction box sizes and misaligned class score meanings in heterogeneous distillation.
[0016] In a possible implementation, the calculation of the intersection over union of the prediction box of the output layer of the student model and the prediction box of the output layer of the teacher model, and the selection of a target prediction box that matches the prediction box of the output layer of the teacher model from the prediction boxes of the output layer of the student model according to the intersection over union include: dividing a feature map corresponding to the input image into a plurality of grids, and calculating a grid to which a center point of each teacher prediction box belongs in the feature map; in the grid to which the center point of each teacher prediction box belongs, calculating an intersection over union of each student prediction box and the teacher prediction box in the grid; and selecting a student prediction box corresponding to the intersection over union with the largest value as the target prediction box.
[0017] In the possible implementation, the center point of each teacher prediction box is determined to belong to a grid, the intersection over union of each teacher prediction box and each student prediction box in the grid to which the center point belongs is calculated, and the most matched student prediction box is selected according to the calculated intersection over union, and the feasibility of the scheme is proposed. In addition, the matching range is limited in the same feature grid, and the case that the offset value in the training process of the student model is limited in an interval can be coped with.
[0018] In a possible implementation, the intersection over union of the prediction box of the student model output layer and the prediction box of the teacher model output layer is calculated, and the target prediction box matched with the prediction box of the corresponding teacher model output layer is selected from the prediction boxes of the plurality of student model output layers according to the intersection over union, and the method comprises: for each teacher prediction box, the intersection over union of the teacher prediction box and each student prediction box in the feature map is calculated; and the student prediction box corresponding to the maximum intersection over union value is selected as the target prediction box.
[0019] In the possible implementation, the most matched student prediction box is selected for each teacher prediction box by calculating the intersection over union of each teacher prediction box and all student prediction boxes in the feature map, the prediction box alignment in the case that the student model prediction box has no interval limitation can be realized, and the feasibility of the scheme is improved.
[0020] In a possible implementation, the output distillation loss is determined, comprising: calculating the output distillation loss based on the product of the output probability distribution of the student model and the foreground confidence and the output probability distribution of the teacher model.
[0021] In the possible implementation, a specific implementation method for determining the output distillation loss is proposed, and the feasibility of the scheme is improved. In addition, by calculating the output distillation loss, the loss function of the student model can be improved based on the output distillation loss in the subsequent process, so that the student model obtained by training has better target detection performance.
[0022] In a possible implementation, the output distillation loss in the output layer is determined, comprising: normalizing the output probability of the student model corresponding to the target probability distribution and the output probability of the teacher model, and determining the output distillation loss based on the normalized result.
[0023] In the possible implementation, another specific implementation method for determining the output distillation loss is proposed, and the feasibility of the scheme is improved. In addition, by calculating the output distillation loss, the loss function of the student model can be improved based on the output distillation loss in the subsequent process, so that the student model obtained by training has better target detection performance.
[0024] In a possible implementation, the loss function of the student model is improved according to the feature distillation loss and the output distillation loss, including: calculating the improved loss function of the student model according to the detection loss of the student model before improvement, the output distillation loss, and the feature distillation loss.
[0025] In this possible implementation, the calculation of the improved loss function of the student model can be implemented to guide the training of the student model based on the improved loss function, so that the student model has better target detection function, and the implementability of the scheme is improved.
[0026] In a second aspect, an embodiment of the present application provides a model training apparatus based on knowledge distillation, which has the function of implementing the model training method based on knowledge distillation in any of the embodiments of the first aspect. The function can be implemented by hardware, or by hardware executing corresponding software. The hardware or software includes one or more modules corresponding to the above functions.
[0027] In a third aspect, the present application provides an electronic device including a memory and a processor. The memory and the processor are coupled. The memory is configured to store computer program code, the computer program code including computer instructions. When the processor executes the computer instructions, the electronic device executes the model training method based on knowledge distillation as described in the first aspect and any possible design of the first aspect.
[0028] In a fourth aspect, the present application provides a computer readable storage medium storing computer instructions, when the computer instructions run on an electronic device, the electronic device executes the model training method based on knowledge distillation as described in the first aspect and any possible design of the first aspect.
[0029] In a fifth aspect, the present application provides a computer program product including computer instructions, when the computer instructions run on an electronic device, the electronic device executes the model training method based on knowledge distillation as described in the first aspect and any possible design of the first aspect.
[0030] The specific description of the second aspect to the fifth aspect and various implementation manners thereof in the present application can refer to the detailed description in the first aspect and various implementation manners thereof; and the beneficial effects of the second aspect to the fifth aspect and various implementation manners thereof can refer to the beneficial effect analysis in the first aspect and various implementation manners thereof, which will not be repeated here.
[0031] These aspects or other aspects of the present application will be more apparent in the following description. BRIEF DESCRIPTION OF DRAWINGS
[0032] Figure 1A schematic diagram of the prior art in which the feature scales of the teacher model and the student model are inconsistent;
[0033] Figure 2 A schematic diagram of an implementation environment involved in a model training method based on knowledge distillation provided by an embodiment of the present application;
[0034] Figure 3 A flowchart of a model training method based on knowledge distillation provided by an embodiment of the present application;
[0035] Figure 4 A schematic diagram of determining a target region of an input image in a teacher model in a model training method based on knowledge distillation provided by an embodiment of the present application;
[0036] Figure 5 A schematic diagram of calculating feature similarity and matching scores based on an attention mechanism in a model training method based on knowledge distillation provided by an embodiment of the present application;
[0037] Figure 6 A schematic diagram of a method of determining a target prediction box in a model training method based on knowledge distillation provided by an embodiment of the present application;
[0038] Figure 7 A schematic diagram of a heterogeneous distillation overall framework of a specific embodiment of a model training method based on knowledge distillation provided by an embodiment of the present application;
[0039] Figure 8 A flowchart of a specific embodiment of a model training method based on knowledge distillation provided by an embodiment of the present application;
[0040] Figure 9 A structural schematic diagram of a target detection model training device based on knowledge distillation provided by an embodiment of the present application;
[0041] Figure 10 A structural schematic diagram of an electronic device provided by an embodiment of the present application. DETAILED DESCRIPTION
[0042] The technical solutions in the embodiments of the present application will be described clearly and completely below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, rather than all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by those of ordinary skill in the art without creative work fall within the scope of protection of the present application.
[0043] The terms "first", "second", etc. are used only for descriptive purposes and should not be construed as implying or suggesting relative importance or an indicated number of the technical features. Thus, the features defined with "first", "second" can include one or more of the features explicitly or implicitly. In the description of the present application, the meaning of "a plurality of" is two or more, unless otherwise specified.
[0044] In the embodiments of the present application, the words "exemplary" or "for example" are used to mean serving as an example, instance, or illustration, in no way limiting. Any embodiment or design scheme described as "exemplary" or "for example" in the embodiments of the present application should not be construed as being more preferred or advantageous than other embodiments or design schemes. Rather, the use of the words "exemplary" or "for example" is intended to present related concepts in a concrete manner.
[0045] Firstly, the technical terms related to the embodiments of the present application are introduced:
[0046] 1. Knowledge distillation
[0047] Knowledge distillation is a common method of model compression. Unlike pruning and quantization, knowledge distillation is to train a small lightweight model by using the supervision information of a large model with better performance, so as to achieve better performance and accuracy. Among them, the large model used for training is the teacher model, and the lightweight model trained is the student model.
[0048] 2. Object detection
[0049] The task of object detection is to find all the targets of interest in the image and determine their categories and positions, which is one of the core problems in the field of computer vision.
[0050] Knowledge distillation can transfer the knowledge of a teacher model with relatively stronger learning ability to a student model with relatively weaker learning ability. Existing knowledge distillation techniques are mostly for the case of teacher model and student model being isomorphic models. However, when the teacher model and the student model are heterogeneous models, for example, the teacher model in knowledge distillation is a two-stage detection network and the student model is a single-stage detection network, there is no complete solution to solve the problems of mismatching of prediction box size, mismatching of class score significance, and mismatching of feature size in heterogeneous distillation.
[0051] Wherein, the inconsistent candidate box size refers to that in the heterogeneous distillation, the prediction box size of the student model and the teacher model is completely different, and cannot be directly matched; the inconsistent category score significance refers to that in the heterogeneous distillation, the category output significance of the teacher model and the student model is inconsistent. Taking the teacher model as a two-stage model and the student model as a single-stage model as an example, the category score output of the teacher model is C+1 dimensions, wherein C is the number of categories, and the extra one dimension is the score of the current region being a background region, which is activated using a softmax function (flexible maximum transmission function, an activation function); the category score output of part of the single-stage student model is C dimensions, and the score of the current region being foreground is output through a conf branch (a confidence branch for distinguishing whether the current region is a foreground region or a background region), which is activated using a sigmoid function (a kind of activation function with exponential function shape, which is most close to biological neurons in physical sense). The difference in activation function means that the functional meanings of the two are different; the inconsistent feature scale refers to that when the heterogeneous distillation target is neck (a series of network layers for mixing and combining image features) features, the downsampling multiple of the teacher model and the downsampling multiple of the student model cannot be aligned, such as shown in Figure 1 In actual problems, the problem of inconsistent feature scale of the teacher model and the student model is more serious. For example, in one possible case, the teacher model neck outputs 5 scale features, and the student model neck only outputs 3 or even 1 scale feature. If the traditional homogeneous feature distillation scheme is followed, at least 2 layers of teacher neck features cannot participate in distillation, and the utilization rate of teacher feature knowledge is low.
[0052] In order to solve the above problems existing in the prior art, the present application proposes a model training method based on knowledge distillation. The implementation mode of the embodiment of the present application will be described in detail below in combination with the drawings:
[0053] Please refer to Figure 2 , which shows an implementation environment schematic diagram of a model training method based on knowledge distillation provided by the embodiment of the present application. As shown in Figure 2 , the implementation environment can include a terminal 210, a terminal 220 and a server 230. The terminal 210 and the terminal 220 are used to deploy a teacher model or a student model, and the server 230 is used for heterogeneous distillation and training the student model based on the heterogeneous distillation process.
[0054] Exemplarily, the terminal 210 is configured to deploy a teacher model, and the terminal 220 is configured to deploy a student model, where the teacher model and the student model are heterogeneous models. In this case, the model training method based on knowledge distillation comprises the following steps: deploying a teacher model for object detection, which has been trained, to the terminal 210, and deploying a student model for object detection to the terminal 220; extracting feature information of the teacher model from the terminal 210, and extracting feature information of the student model from the terminal 220; when the feature information of the student model is extracted, matching the feature information of the student model with the feature information of the teacher model, and calculating a feature distillation loss; extracting output information of the teacher model from the terminal 210, and extracting output information of the student model from the terminal 220; matching the output information of the student model with the output information of the teacher model, and calculating an output distillation loss; improving a loss function of the student model based on the calculated feature distillation loss function and output distillation loss function, and guiding the training of the student model based on the improved loss function, so that the student model obtained through the training has similar object detection function to the teacher model.
[0055] It should be noted that the terminal 210 and the terminal 220 in the embodiments of the present application can be a mobile phone, a tablet computer, a notebook computer, and other devices with data processing capabilities, and the specific form of the terminal 210 and the terminal 220 is not specially limited in the embodiments of the present application.
[0056] The student model obtained through the training of the embodiments of the present application can be used in various scenes in the field of computer vision application, or a processing system based on a neural network model deployed on an edge device (such as a mobile phone, a wearable device, a computing node, etc.), or an application scene in which the neural network model needs to be compressed due to limited resources and latency requirements.
[0057] Exemplarily, the terminal 220 can be a smart phone, and the student model not only retains similar object detection function to the teacher model, but also can be deployed on a terminal with limited resources due to its simple structure. Therefore, the student model can be deployed on the smart phone terminal 220, and the user can use the phone to automatically capture faces, animals and other targets, so as to help the phone to automatically focus, beautify, etc. when taking pictures, thereby providing better user experience for the user. In addition, the terminal 220 can also be used in the application scene of automatic driving scene segmentation. After the camera of the automatic driving vehicle captures the road picture, the picture needs to be segmented to separate different objects such as road surface, roadbed, vehicle and pedestrian, so as to keep the vehicle driving in the correct area. The embodiments of the present application can be adjusted and improved according to specific application environment, which is not specifically limited here.
[0058] The following will be described in combination with Figure 3The flowchart shown provides a detailed description of the model training method based on knowledge distillation provided by the embodiments of the present application, such as Figure 3 As shown, the method can include S301-S305.
[0059] S301, obtaining a trained teacher model.
[0060] The teacher model can be used for target detection of the input image.
[0061] The method proposed in the example embodiment is used for training of a target detection model using knowledge distillation when the teacher model and the student model are heterogeneous models, for example, the teacher model involved in knowledge distillation is a two-stage detection network, and the student model is a single-stage detection network.
[0062] The teacher model is a large model that provides more powerful supervision information in knowledge distillation, and the student model is a small model that learns the knowledge of the teacher model. For example, the teacher model can be a model that optimizes anchor boxes in two or more stages, such as Faster RCNN, Cascade RCNN, etc. The student model can be a model that only optimizes anchor boxes in a single stage, such as RetinaNet, YOLO, etc. Because the optimization stages of the teacher model and the student model are different, the anchor box sizes are also different, and direct matching cannot be performed during knowledge distillation. The solution proposed in this embodiment is used to solve this heterogeneous distillation problem.
[0063] It should be noted that the above case is only an example, and other cases where the teacher model and the student model are heterogeneous models also fall within the protection scope of the example embodiment.
[0064] In the example embodiment, the teacher model is a pre-trained model, which has strong target detection capability but relatively complex structure and relatively high requirements for deployment resources. The input image refers to image data input to the teacher model. The teacher model used for target detection of the input image can be understood as the teacher model classifying and positioning objects in the image data of the input model. For example, the input image can be a photo containing multiple objects such as people, dogs, and trees. The photo is input to the teacher model, which can detect whether the target object in the image is a person, a dog, or a tree, and can locate the position of the person in the photo. It should be noted that the above scenario is only an example, and the example embodiment does not have special limitations.
[0065] S302, align the feature scales of the student model feature layer and the teacher model feature layer, and determine the feature distillation loss at the feature layer.
[0066] The characteristic distillation loss represents a difference between the teacher model and the student model in the characteristic layer after the characteristic scales are aligned.
[0067] In the example embodiment, the alignment of the characteristic scales of the student model characteristic layer and the characteristic scales of the teacher model characteristic layer can be implemented as steps 1 to 3 as follows:
[0068] Step 1: Obtain a teacher characteristic set of the teacher model.
[0069] The teacher characteristics in the teacher characteristic set are characteristics of a target region in a first image input to the teacher model, and the target region is determined based on an anchor box. The characteristics of the target region include foreground region characteristics and background region characteristics of the target region.
[0070] Optionally, a target region can be determined in the input image based on the anchor box, and a target region characteristic corresponding to the target region is obtained as the teacher characteristic set. The target region characteristic includes both the foreground characteristics of the target region and the background characteristics of the target region, as shown in FIG. 4. Figure 4 As shown in FIG. 4, 401 is an anchor box of an image, and the region determined by the anchor box is the target region. The foreground characteristics and the background characteristics in the region are both target region characteristics.
[0071] It should be noted that the above scenario is only an example, and the protection scope of the example embodiment is not limited thereto. The target region can also be other regions in the image.
[0072] Step 2: Obtain a student characteristic set in the student model.
[0073] The student characteristics in the student characteristic set are characteristics of the target region in the student model.
[0074] Optionally, the position of the target region of the input image (e.g., the first image) of the teacher model determined by the anchor box can be determined, and the region characteristics of the student model in the corresponding position of the image (e.g., the first image) can be obtained when the image is input to the student model, to obtain the student characteristic set.
[0075] Step 3: Convert the characteristic scales in the teacher characteristic set and the student characteristic set to the same dimension.
[0076] Exemplarily, the conversion of the feature scales of the teacher feature set and the student feature set to the same dimension can be implemented as follows: sequentially performing scale transformation and channel dimension transformation on the student feature set, so that the feature scales of the teacher model and the student model are converted to the same dimension. Specifically, the conversion process can be implemented by a dimension conversion module, which can include RoIAlign (used to convert a feature map of an arbitrary size of a target region into a small feature map with a fixed size) and an Adaptive Layer.
[0077] The determination of the feature distillation loss at the feature layer can be implemented as follows: calculating the feature similarity between each teacher feature in the teacher feature set and each student feature in the student feature set, and normalizing to obtain a corresponding matching score; and calculating the feature distillation loss according to each matching score and the distillation loss corresponding to each matching score.
[0078] Exemplarily, the calculation of the feature similarity between each teacher feature in the teacher feature set and each student feature in the student feature set, and the normalization to obtain a corresponding matching score can be implemented in a manner based on an attention mechanism or a cosine similarity.
[0079] Specifically, the manner based on the attention mechanism can be implemented as shown in the flowchart of FIG. 1B. Figure 5 Specifically, the manner based on the attention mechanism can be implemented as shown in the flowchart of FIG. 1B.
[0080] Alternatively, the manner based on the cosine similarity algorithm can be implemented as follows: after the teacher features and the student features pass through avg pooling (average filtering convolution), the cosine similarity between the teacher features and the student features is calculated, and the matching score a is obtained by normalization through a softmax function. The above matching process is to match each pair of features in the teacher feature set and the student feature set to obtain the corresponding matching score a.
[0081] Exemplarily, the calculation of the feature distillation loss according to each matching score and the distillation loss corresponding to each matching score can be calculated based on the following formula:
[0082] L1=
[0083] wherein, f iand f j respectively represent the single-layer teacher feature and student feature, and a represents the matching score calculated by the above process. LMSE(, ) represents the MSE loss function.
[0084] S303, align the prediction box and the corresponding probability distribution of the output layer of the student model with the prediction box and the corresponding probability distribution of the output layer of the teacher model, and determine the output distillation loss of the output layer.
[0085] The output distillation loss represents the difference between the teacher model and the student model in the output layer after the prediction box and the corresponding probability distribution are aligned, and the probability distribution is the probability distribution output by the prediction box.
[0086] In the example embodiment, the step first obtains the prediction box and the corresponding probability distribution of the output layer of the teacher model, then aligns the prediction box and the corresponding probability distribution of the output layer of the student model with the prediction box and the corresponding probability distribution of the output layer of the teacher model, and calculates the output distillation loss to complete the knowledge distillation of the output layer.
[0087] In the example embodiment, the teacher prediction box and the student prediction box are anchor boxes used for positioning in the input image, the probability distribution of the teacher model can be the class probability distribution output by the corresponding teacher prediction box, and the probability distribution of the student model can be the class probability distribution output by the corresponding student prediction box.
[0088] For example, the alignment of the prediction box and the corresponding probability distribution of the output layer of the student model with the prediction box and the corresponding probability distribution of the output layer of the teacher model can be implemented as follows: calculate the intersection over union of the student prediction box and the teacher prediction box, and select a target prediction box matching the corresponding teacher prediction box from the multiple student prediction boxes according to the intersection over union; obtain the target probability distribution corresponding to the target prediction box; and align the meaning of the target probability distribution with the probability distribution corresponding to the matching teacher prediction box.
[0089] In one possible implementation, the calculation of the intersection over union of the student prediction box and the teacher prediction box, and the selection of a target prediction box matching the corresponding teacher prediction box from the multiple student prediction boxes according to the intersection over union can be implemented as follows:
[0090] Divide the feature map corresponding to the input image into multiple grids; calculate the grid to which the center point of each teacher prediction box belongs in the feature map; in the grid to which the center point of each teacher prediction box belongs, calculate the intersection over union of each student prediction box and the teacher prediction box in the grid; and select the student prediction box corresponding to the largest intersection over union as the target prediction box.
[0091] For example, the alignment of the prediction box and the corresponding probability distribution of the output layer of the student model with the prediction box and the corresponding probability distribution of the output layer of the teacher model can be implemented as follows: Figure 6 For example, the alignment of the prediction box and the corresponding probability distribution of the output layer of the student model with the prediction box and the corresponding probability distribution of the output layer of the teacher model can be implemented as follows:Figure 6 As shown, the input image is divided into multiple grids, and for each teacher prediction box, the grid to which the center point of the teacher prediction box belongs in the image is calculated, and it is assumed that the grid to which the current teacher prediction box belongs is Figure 6 Grid 601 in the image, all student prediction boxes in the grid 601 are obtained, and the intersection over union of each student prediction box and the current teacher box is calculated, and the student prediction box corresponding to the maximum value of the intersection over union is selected as the target prediction box corresponding to the current teacher prediction box. The above process is repeated for all other teacher prediction boxes to obtain the corresponding target prediction box.
[0092] It should be noted that the above possible implementation limits the matching range to the same feature grid, and can cope with the case that the offset value of the prediction box in the training process of the student model has an interval limit. In this case, there is no interval limit in the supervision value provided by the teacher, but the student model has an upper limit on the interval and cannot be completely fitted.
[0093] In another possible implementation, the above calculation of the intersection over union of the student prediction box and the teacher prediction box, and the selection of the target prediction box matching the corresponding teacher prediction box from the multiple student prediction boxes according to the intersection over union can be implemented as follows:
[0094] For each teacher prediction box, the intersection over union of the teacher prediction box and each student prediction box in the feature map is calculated; and the student prediction box corresponding to the maximum value of the intersection over union is selected as the target prediction box.
[0095] In this possible implementation, the feature map of the input image is taken as a whole without grid division. For each teacher prediction box, all student prediction boxes in the feature map are obtained, and the intersection over union of the teacher prediction box and all student prediction boxes is calculated, and the student prediction box corresponding to the maximum value of the intersection over union is selected as the target prediction box. It should be noted that the above possible implementation is applicable to the case that the offset value of the prediction box of the student model has no interval limit.
[0096] Further, in some embodiments, the above determination of the output distillation loss can be implemented as follows: multiplying the target probability distribution by the foreground confidence of the student model to align the meaning of the target probability distribution with the probability distribution corresponding to the matching teacher prediction box; calculating the output distillation loss based on the product of the target probability distribution and the foreground confidence, the output probability distribution of the teacher model and the student model. The specific calculation formula of the output distillation loss can be as follows:
[0097]
[0098] Wherein, sigmoid(q) is the output probability distribution of the student model, softmax(p) is the output probability distribution of the teacher model, the operations of the softmax and sigmoid functions are consistent with the network training process, conf is the foreground confidence of the student model, which is used to represent the probability that the region determined by the current prediction frame is a foreground region, and the loss function h() can be in the form of KL loss, ce loss, etc.
[0099] In addition, in another embodiment, the determination of the output distillation loss can also be implemented as follows: the output probability of the student model corresponding to the target probability distribution and the output probability of the teacher model are normalized to align the meanings of the target probability distribution and the probability distribution corresponding to the matching teacher prediction frame; the output distillation loss is calculated based on the output probability distribution of the student model corresponding to the target probability distribution and the output probability distribution of the teacher model, and the specific calculation formula of the output distillation loss can be as follows:
[0100] L2= h(softmax(p), softmax(q))
[0101] Wherein, softmax(p) is the output probability distribution of the student model, softmax(q) is the output probability distribution of the teacher model, softmax and sigmoid are normalization functions, and the loss function h() can be in the form of KL loss, ce loss, etc.
[0102] S304, improving the loss function of the student model according to the feature distillation loss and the output distillation loss.
[0103] In the example embodiment, this step is used to improve the loss function of the student model, and the training of the student model is guided based on the improved loss function to obtain a student model with similar target detection function as the teacher model. The process can be implemented as follows: the improved loss function of the student model is calculated according to the original detection loss of the student model, the feature distillation loss L1 and the output distillation loss L2 calculated above.
[0104] For example, the calculation of the improved loss function can be implemented as follows: multiplying the output distillation loss by the corresponding output weight parameter to obtain a first product; multiplying the feature distillation loss by the corresponding feature weight parameter to obtain a second product; adding the original detection loss, the first product and the second product to obtain the improved loss function, and the specific calculation formula of the improved loss function is as follows:
[0105] L=Ldet + β•L1+γ•L2
[0106] Wherein, Ldet is the original detection loss of the student model, L1 and L2 are feature heterogeneous distillation loss and output heterogeneous distillation loss respectively, and β and γ are loss weight hyperparameters corresponding to the feature heterogeneous distillation loss and the output heterogeneous distillation loss respectively.
[0107] S305, training the student model based on the improved loss function to obtain a trained model.
[0108] Below, combined with Figure 7 With Figure 8 The specific application scenario shown in the figure, the above-mentioned model training method based on knowledge distillation is described in detail:
[0109] Figure 7 The heterogeneous distillation overall framework of the above-mentioned model training method based on knowledge distillation is applied to the case where the teacher model and the student model are heterogeneous models, and the target detection model is trained by means of knowledge distillation. The framework includes a two-stage teacher model and a single-stage student model, the corresponding feature layer information and output layer information are extracted from the teacher model and the student model respectively, and the feature heterogeneous distillation and the output heterogeneous distillation are performed to improve the student model.
[0110] For Figure 7 The specific implementation process is shown in Figure 8 , including the following steps:
[0111] S801, obtaining a trained teacher model, which is used for target detection on input images.
[0112] This step is used to obtain a pre-trained teacher model, which is a large model with stronger functions in knowledge distillation and is used to provide supervision information to realize target detection.
[0113] S802, feature heterogeneous distillation of the teacher model and the student model.
[0114] This step is used to perform feature heterogeneous distillation on the teacher model and the student model, and to calculate the feature heterogeneous distillation loss. The specific implementation is as follows:
[0115] S8021, obtaining the target region features of the above-mentioned teacher model, which contains both the foreground region and the background region, to obtain the teacher feature set FeatT corresponding to the teacher model.
[0116] S8022: obtaining the region features corresponding to the target region of the teacher model in the features of each layer of the student model to obtain the student feature set FeatS.
[0117] S8023: transforming the scales of the features in the teacher feature set and the student feature set to the same dimension through the dimension conversion module.
[0118] Specifically, the aforementioned dimension transformation module may include RoIAlign and Adaptive Layer (1x1 conv). This step can be implemented by performing scale transformation and channel dimension transformation on FeatureS sequentially through the dimension transformation module, thereby transforming the feature scale of the teacher feature set and the student feature set to the same dimension.
[0119] S8024: Perform many-to-many matching on the transformed feature set, and use the matching score to weight the loss.
[0120] Implementation Method 1: An attention mechanism can be used to calculate the matching score between features. The teacher feature set is sequentially processed through average pooling (average filtering convolution) and a fully connected query layer. The student features, whose target regions are located in the corresponding regions, are processed through average pooling (average filtering convolution) and a fully connected keylayer. The similarity between the teacher and student features is then calculated and normalized using a softmax function to obtain the matching score α.
[0121] Implementation Method 2: Calculate the matching score using cosine similarity. After average pooling, the cosine similarity between the teacher and student feature sets is calculated, and finally, the matching score α is obtained through softmax normalization.
[0122] S8025: The matching score is multiplied by the corresponding feature distillation loss to obtain the final feature heterogeneous distillation loss.
[0123] L1=
[0124] Where fi and fj represent single-layer teacher and student features, respectively, and α represents the matching score calculated in the above process. LMSE(,) represents the MSE loss function.
[0125] S803, Perform heterogeneous distillation on the output of the teacher model and the student model.
[0126] This step is used to perform output heterogeneous distillation on the teacher and student models and calculate the output heterogeneous distillation loss. The specific implementation is as follows:
[0127] S8031: Obtain the prediction boxes for the teacher model and the student model.
[0128] S8032: Perform position matching on the predicted bounding boxes of the teacher model and the student model. The matching method is as follows: calculate the intersection-union ratio (IU) of the predicted bounding boxes of the student model and the teacher model, and select the predicted bounding box with the largest IU for online matching during training. For example, the specific implementation can be as follows:
[0129] Implementation 1: The feature map is divided into a grid of SxS (S is a positive integer); the grid to which each teacher prediction box center point belongs is calculated; the intersection over union is calculated between each teacher prediction box and all student prediction boxes in the grid to which the center point belongs; the student prediction box with the largest intersection over union is selected as the matching pair.
[0130] This implementation limits the matching range to the same feature grid, which can handle the case where the student model's prediction box offset value has an interval limit during training. In this case, the teacher provides supervision values without interval limits, while the student model has an upper limit and cannot be completely fitted.
[0131] Implementation 2: The intersection over union is calculated between each teacher prediction box and all student prediction boxes in the feature map; the student prediction box with the largest intersection over union is selected as the matching pair. This implementation is suitable for cases where the student model's prediction box offset value has no interval limit.
[0132] S8033: Obtain the probability distribution corresponding to the matched teacher model prediction box and student model prediction box.
[0133] S8034: Align the probability distribution of the teacher model prediction box with the probability distribution of the student model prediction box, and calculate the output distillation loss.
[0134] Implementation 1: The specific output distillation loss form is shown by the following formula:
[0135]
[0136] where q is the output probability distribution of the student model, p is the output probability distribution of the teacher model, the softmax and sigmoid functions are operated in accordance with the network training process, and conf is the foreground confidence of the student model, which is used to represent the probability that the current prediction box determines the region as a foreground region. Multiplying the student's foreground confidence score conf with the classification probability distribution can align the probability distribution meaning, and the loss function h() can be in the form of KL loss, ce loss, etc.
[0137] Implementation 2:
[0138] L2 = h(softmax(p), softmax(q))
[0139] The normalized forms of the output probabilities of the student model and the teacher model are consistent, the unified normalization obtains the output probability distribution with consistent meanings, and the distillation loss is calculated through the loss function h(). Wherein, q is the output probability distribution of the student model, p is the output probability distribution of the teacher model, softmax and sigmoid are normalization functions, and the loss function h() can be in the form of KL loss, ce loss, etc.
[0140] S804, improve the loss function of the student model, and train the student model based on the improved loss function.
[0141] The specific calculation formula of the improved loss function is as follows:
[0142] L = Ldet + β • L1 + γ • L2
[0143] Wherein, Ldet is the original detection loss of the student model, L1 and L2 are feature heterogeneous distillation loss and output heterogeneous distillation loss respectively, and β and γ are loss weight hyperparameters corresponding to the feature heterogeneous distillation loss and the output heterogeneous distillation loss respectively.
[0144] The example embodiment provides a complete heterogeneous distillation scheme, which can use a two-stage detection network with better performance as a teacher model in the training process of the target detection model, thereby improving the performance upper limit of distillation and expanding the selection range of the teacher and the application range of the distillation algorithm. In addition, the example scheme has been verified to be effective in the Adas (Advanced Driver Assistance System, advanced driving assistance system) business scenario and the public data set VOC.
[0145] The above mainly introduces the scheme provided by the embodiments of the application from the method aspect. In order to realize the above functions, it contains the hardware structure and / or software module corresponding to the execution of each function. The technical person skilled in the art should easily realize that the units and algorithm steps of each example described in combination with the embodiments disclosed in the present text can be realized in the form of hardware or a combination of hardware and computer software. Whether a certain function is executed in hardware or computer software driven hardware depends on the specific application and design constraints of the technical scheme. The professional technical person skilled in the art can use different methods to realize the described functions for each specific application, but such implementation should not be considered beyond the scope of the present application.
[0146] The embodiments of the present application also provide a target detection model training device based on knowledge distillation, wherein the teacher model and the student model are heterogeneous models, for example, the teacher model involved in knowledge distillation is a two-stage detection network, and the student model is a single-stage detection network. As Figure 9As shown, a structural schematic diagram of a model training device 900 based on knowledge distillation is provided. The device 900 can include a teacher model obtaining module 901, a feature distillation module 902, an output distillation module 903, and a student model training module 904, wherein:
[0147] The teacher model obtaining module 901 can be used to obtain the trained teacher model.
[0148] The feature distillation module 902 can be used to align the feature scales of the student model feature layers and the feature scales of the teacher model feature layers, and determine a feature distillation loss, which represents the difference between the teacher model and the student model in the feature layers after the alignment of the feature scales.
[0149] The output distillation module 903 can be used to align the prediction boxes and the corresponding probability distributions of the student model output layers with the prediction boxes and the corresponding probability distributions of the teacher model output layers, and determine an output distillation loss, which represents the difference between the teacher model and the student model in the output layers after the alignment of the prediction boxes and the corresponding probability distributions. The probability distribution is the probability distribution of the output of the prediction box.
[0150] The student model training module 904 can be used to improve the loss function of the student model according to the feature distillation loss and the output distillation loss, and train the student model based on the improved loss function to obtain a trained model.
[0151] In a possible implementation, the feature distillation module is specifically configured to: obtain a teacher feature set of the teacher model; wherein the teacher features in the teacher feature set are features of a target region in a first image input to the teacher model, the target region is a region determined based on an anchor box, and the features of the target region include foreground region features and background region features of the target region; obtain a student feature set in the features of each layer of the student model; wherein the student features in the student feature set are features of the target region in the student model; and convert the feature scales in the teacher feature set and the student feature set to the same dimension.
[0152] Specifically, the student feature set obtained in the features of each layer of the student model includes: extracting student region features of each layer based on the target region in the first image input to the student model to obtain the student feature set.
[0153] Specifically, the feature distillation loss in the feature layers is determined by: calculating the feature similarity between each teacher feature in the teacher feature set and each student feature in the student feature set based on an attention mechanism or a cosine similarity algorithm, and normalizing to obtain a corresponding matching score; and calculating the feature distillation loss according to each matching score and a distillation loss corresponding to each matching score.
[0154] In a possible implementation, the output distillation module is specifically configured to: calculate an intersection over union of a prediction box of the student model output layer and a prediction box of the teacher model output layer, and select a target prediction box matching the prediction box of the corresponding teacher model output layer from the prediction boxes of the plurality of student model output layers according to the intersection over union; obtain a target probability distribution corresponding to the target prediction box, the target probability distribution being an output probability distribution of the target prediction box; and align the target probability distribution with a probability distribution corresponding to the prediction box of the teacher model output layer.
[0155] In a possible implementation, the calculation of the intersection over union of the prediction box of the student model output layer and the prediction box of the teacher model output layer, and the selection of the target prediction box matching the prediction box of the corresponding teacher model output layer from the prediction boxes of the plurality of student model output layers according to the intersection over union, includes: dividing a feature map corresponding to the input image into a plurality of grids, and calculating a grid to which a center point of each teacher prediction box belongs in the feature map; in the grid to which the center point of each teacher prediction box belongs, calculating an intersection over union of each student prediction box and the teacher prediction box in the grid; and selecting a student prediction box corresponding to a largest intersection over union as the target prediction box.
[0156] In another possible implementation, the calculation of the intersection over union of the prediction box of the student model output layer and the prediction box of the teacher model output layer, and the selection of the target prediction box matching the prediction box of the corresponding teacher model output layer from the prediction boxes of the plurality of student model output layers according to the intersection over union, includes: for each teacher prediction box, calculating an intersection over union of the teacher prediction box and each student prediction box in the feature map; and selecting a student prediction box corresponding to a largest intersection over union as the target prediction box.
[0157] In a possible implementation, the determination of the output distillation loss includes: calculating the output distillation loss based on a product of the output probability distribution of the student model and the foreground confidence and the output probability distribution of the teacher model.
[0158] In another possible implementation, the determination of the output distillation loss includes: normalizing the output probability of the student model and the output probability of the teacher model corresponding to the target probability distribution to align the target probability distribution with the probability distribution corresponding to the matching teacher prediction box in meaning; and calculating the output distillation loss based on the output probability distribution of the student model and the output probability distribution of the teacher model corresponding to the target probability distribution.
[0159] In a possible implementation, the improvement of the loss function of the student model according to the feature distillation loss and the output distillation loss includes: calculating the improved loss function of the student model according to the detection loss of the student model before improvement, the output distillation loss, and the feature distillation loss.
[0160] Of course, the model training apparatus 900 based on knowledge distillation provided by the embodiments of the present application includes but is not limited to the above modules.
[0161] Another embodiment of the present application also provides an electronic device. As shown in Figure 10 The electronic device 1000 includes a memory 1001 and a processor 1002; the memory 1001 and the processor 1002 are coupled; the memory 1001 is configured to store computer program code, and the computer program code includes computer instructions. When the processor 1002 executes the computer instructions, the electronic device 1000 performs each step of the method flow shown in the above method embodiment.
[0162] In actual implementation, the teacher model obtaining module 901, the feature distillation module 902, the output distillation module 903, and the student model training module 904 can be implemented by Figure 10 The specific execution process can refer to the description of the above method part, and will not be described here.
[0163] Another embodiment of the present application also provides a computer readable storage medium, and the computer readable storage medium stores computer instructions. When the computer instructions run on the electronic device, the electronic device performs each step of the method flow shown in the above method embodiment.
[0164] In another embodiment of the present application, a computer program product is also provided, and the computer program product includes computer instructions. When the computer instructions run on the electronic device, the electronic device performs each step of the method flow shown in the above method embodiment.
[0165] In the above embodiments, all or part of the embodiments can be implemented by software, hardware, firmware or any combination thereof. When implemented by software, all or part of the embodiments can be implemented in the form of a computer program product. The computer program product includes one or more computer instructions. When the computer executes the computer instructions, all or part of the processes or functions according to the embodiments of the present application are generated. The computer can be a general purpose computer, a special purpose computer, a computer network or other programmable device. The computer instructions can be stored in a computer readable storage medium or transmitted from one computer readable storage medium to another computer readable storage medium, for example, the computer instructions can be transmitted from one website, computer, server or data center to another website, computer, server or data center through wired (for example, coaxial cable, optical fiber, digital subscriber line (DSL)) or wireless (for example, infrared, wireless, microwave, etc.). The computer readable storage medium can be any available medium that the computer can access or include one or more data storage devices that can be integrated with the medium, such as servers, data centers, etc. The available medium can be a magnetic medium (for example, floppy disk, hard disk, magnetic tape), an optical medium (for example, DVD) and the like.
[0166] The above is only a specific embodiment of the present application. Those skilled in the art can think of changes or replacements based on the specific embodiments provided by the present application, which should be covered within the protection scope of the present application.
Claims
1. A model training method based on knowledge distillation, characterized in that, The teacher model and the student model involved in the knowledge distillation are heterogeneous models, and the method comprises: obtaining a trained teacher model, the teacher model being used for target detection on an input image; aligning the feature scales of the feature layers of the student model and the feature scales of the feature layers of the teacher model, and determining a feature distillation loss, the feature distillation loss representing the difference between the feature layers of the teacher model and the student model after the alignment of the feature scales; aligning the prediction boxes and the corresponding probability distributions of the output layers of the student model and the prediction boxes and the corresponding probability distributions of the output layers of the teacher model, and determining an output distillation loss, the output distillation loss representing the difference between the output layers of the teacher model and the student model after the alignment of the prediction boxes and the corresponding probability distributions, the probability distribution being the probability distribution of the output of the prediction box; improving the loss function of the student model according to the feature distillation loss and the output distillation loss; training the student model based on the improved loss function to obtain a trained model.
2. The method of claim 1, wherein, The alignment of the feature scales of the feature layers of the student model and the feature scales of the feature layers of the teacher model comprises: obtaining a teacher feature set of the teacher model; wherein the teacher features in the teacher feature set are features of a target region in a first image input to the teacher model, the target region being a region determined based on an anchor box, and the features of the target region comprising foreground region features and background region features of the target region; obtaining a student feature set from the feature layers of the student model; wherein the student features in the student feature set are features of the target region in the student model; converting the feature scales in the teacher feature set and the student feature set to the same dimension.
3. The method of claim 2, wherein, The obtaining of the student feature set from the feature layers of the student model comprises: extracting student region features of the layers based on the target region in the first image input to the student model to obtain the student feature set.
4. The method according to claim 2 or 3, characterized in that, The determination of the feature distillation loss comprises: calculating feature similarities between each teacher feature in the teacher feature set and each student feature in the student feature set based on an attention mechanism or a cosine similarity algorithm, and determining the feature distillation loss according to the feature similarities.
5. The method of claim 1, wherein, The alignment of the prediction boxes and the corresponding probability distributions of the output layers of the student model and the prediction boxes and the corresponding probability distributions of the output layers of the teacher model comprises: calculating intersection over union of the prediction boxes of the output layers of the student model and the prediction boxes of the output layers of the teacher model, and selecting a target prediction box matching a corresponding prediction box of the output layers of the teacher model from the prediction boxes of the output layers of the student model according to the intersection over union; obtaining a target probability distribution corresponding to the target prediction box, the target probability distribution being the probability distribution of the output of the target prediction box; aligning the target probability distribution with the probability distribution corresponding to the prediction box of the output layers of the teacher model.
6. The method of claim 5, wherein, The calculating the intersection over union of the prediction box of the student model output layer and the prediction box of the teacher model output layer, and selecting a target prediction box matching the prediction box of the corresponding teacher model output layer from the prediction boxes of the student model output layer according to the intersection over union, comprises: dividing a feature map corresponding to an input image into a plurality of grids, and calculating the grid to which the center point of the prediction box of each teacher model output layer belongs in the feature map; in the grid to which the center point of each prediction box of the teacher model output layer belongs, calculating the intersection over union of the prediction box of each student model output layer in the grid and the prediction box of the teacher model output layer; selecting the prediction box of the student model output layer corresponding to the intersection over union with the maximum value as the target prediction box.
7. The method of claim 5, wherein, The calculating the intersection over union of the prediction box of the student model output layer and the prediction box of the teacher model output layer, and selecting a target prediction box matching the prediction box of the corresponding teacher model output layer from the prediction boxes of the student model output layer according to the intersection over union, comprises: for each prediction box of the teacher model output layer, calculating the intersection over union of the prediction box of the teacher model output layer and each prediction box of the student model output layer in the feature map corresponding to the input image; selecting the prediction box of the student model output layer corresponding to the intersection over union with the maximum value as the target prediction box.
8. The method of claim 5, wherein, The determining the output distillation loss comprises: calculating the output distillation loss based on the product of the output probability distribution of the student model and the foreground confidence and the output probability distribution of the teacher model.
9. The method of claim 5, wherein, The determining the output distillation loss comprises: normalizing the output probability of the student model and the output probability of the teacher model, and determining the output distillation loss based on the normalized result.
10. The method according to any one of claims 1-3 or 5-9, characterized in that, The improving the loss function of the student model according to the feature distillation loss and the output distillation loss comprises: calculating the improved loss function of the student model according to the detection loss of the student model before improvement, the output distillation loss and the feature distillation loss. 11.A model training apparatus based on knowledge distillation, characterized in that, The teacher model and the student model involved in knowledge distillation are heterogeneous models, and the device comprises: a teacher model acquisition module configured to acquire a trained teacher model, the teacher model being configured to perform target detection on an input image; a feature distillation module configured to align the feature scales of the feature layers of the student model and the feature layers of the teacher model, and determine a feature distillation loss, the feature distillation loss representing the difference between the feature layers of the teacher model and the student model after the alignment of the feature scales; an output distillation module configured to align the prediction boxes and the corresponding probability distributions of the output layers of the student model and the prediction boxes and the corresponding probability distributions of the output layers of the teacher model, and determine an output distillation loss, the output distillation loss representing the difference between the output layers of the teacher model and the student model after the alignment of the prediction boxes and the corresponding probability distributions, the probability distribution being the output probability distribution of the prediction box. The student model training module is configured to improve a loss function of the student model according to the feature distillation loss and the output distillation loss, and train the student model based on the improved loss function to obtain a trained model.
12. An electronic device, comprising: The electronic device includes a memory and a processor; the memory and the processor are coupled; the memory is used to store computer program code, the computer program code includes computer instructions; When the processor executes the computer instructions, the electronic device executes the model training method based on knowledge distillation as any one of claims 1-10.
13. A computer-readable storage medium, characterized in that, The computer readable storage medium stores computer instructions, when the computer instructions run on the electronic device, the electronic device executes the model training method based on knowledge distillation as any one of claims 1-10.
Citation Information
Patent Citations
Image category identification method and device based on model distillation, storage medium and terminal
CN113408570A
Model training method and system and computer storage medium
CN114897160A