Model compression method and system based on knowledge distillation

Abstract

The invention provides a knowledge distillation-based model compression method and system, and the method comprises the steps: inheriting a soft label of a maturely trained teacher model, defining a distillation loss function of a student model through the soft label, enabling the student model to be compressed, generating a new sample extension training set through random sample extraction and mixing, and achieving the compression of the student model. The problem of insufficient data is relieved, meanwhile, the robustness of the model to noise and distribution offset is improved, through dynamic weight adjustment, a student model is made to quickly fit teacher model knowledge, autonomous optimization is conducted in the later training period, and over-fitting of soft labels is avoided. On a test set in which the proportion of unlabeled data is 30%, on the basis that the accuracy of using a real label is 78.3%, the accuracy is improved by 11.2% by using a pseudo target, deployment is performed by using a maturely trained student model, middle layer feature alignment is omitted, the occupation of a GPU video memory is reduced from 3.2 GB to 1.8 GB, and the deployment requirement of edge equipment is met.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of model compression, and particularly relates to a model compression method and system based on knowledge distillation. BACKGROUND

[0002] Knowledge distillation is a model compression technology that guides student models to learn through teacher models. The method of knowledge distillation based on group relationship proposed by Nantong University optimizes the performance of student models through clustering and maximum mean difference, but it depends on complex intermediate feature alignment and has high calculation cost. The heterogeneous knowledge distillation method based on the Transformer attention mechanism proposed by Space Science and Technology Corporation improves the performance of cross-modal tasks, but does not solve the generalization problem of lightweight models in untrained environments.

[0003] The current mainstream method has the following limitations: (1) Complex feature alignment: existing methods need to additionally calculate the similarity of intermediate layer features (such as attention weights or clustering relationships), increasing the training overhead.

[0004] (2) Insufficient performance in untrained environments: the accuracy of lightweight models significantly decreases in scenarios where they have not participated in training. SUMMARY

[0005] The main problem solved by the application is how to solve the generalization problem of lightweight models in untrained environments and improve the accuracy of models in unknown scenarios. The application provides a model compression method and system based on knowledge distillation.

[0006] To solve the above technical problems, the technical solution adopted is: A model compression method based on knowledge distillation, comprising the following steps: Step 1: train the teacher model using the training set until convergence, obtaining soft labels corresponding to each training sample; Step 2: initialize the student model and define the distillation loss function of the student model according to the soft labels; Step 3: randomly extract two samples from the training set and generate a new mixed sample according to a certain mixing ratio. Repeat this step to generate multiple mixed samples; Step 4: mix the mixed samples and the training set according to a pre-set ratio to enhance the training set; Step 5: generate pseudo labels for all mixed samples in the enhanced training set, the pseudo labels being generated by the teacher model; Step 6: use the enhanced training data set to jointly train the student model, obtaining a lightweight student model for deployment.

[0007] Furthermore, the teacher model is trained using the cross entropy method, and the soft labels corresponding to each sample are obtained as follows:

[0008]

[0009] in, is the i-th probability value of the teacher model output probability, is the i-th probability value of the original probability value of the teacher model output layer, is the true label index of the i-th probability, and N is the number of output probabilities of the teacher model.

[0010] Furthermore, the distillation loss function of the student model is defined according to the soft labels for:

[0011]

[0012]

[0013] in, For supervision of soft labels Divergence loss value, is the temperature parameter used to smooth the probability distribution; Output all original probability values ​​of the student model layer samples; All original probability values ​​of the teacher model output layer; To dynamically adjust the weight coefficient, As the training process decreases linearly, the student model gradually transitions from relying on the teacher's knowledge to independent learning. Indicates the number of times the complete training data set is traversed in a single pass, Represents the total number of training rounds.

[0014] Furthermore, a method for generating a new mixed sample according to a certain mixing ratio is:

[0015]

[0016] Among them, among them, and are two input samples randomly selected from the training data, and The samples and The corresponding true label, is the mixing ratio.

[0017] Further, the cross-entropy loss generated by the enhancement strategy for the newly generated mixed sample is: ; wherein, is the cross-entropy loss function.

[0018] Further, when the student model is jointly trained using the enhanced training set: the enhanced training set sample is input into the student model, the original sample in the training set is trained using the distillation loss function, and the mixed sample is trained using the cross-entropy loss function; the total loss function of the joint training is:

[0019] wherein, is the cross-entropy loss weight.

[0020] The application also provides a model compression system based on knowledge distillation, which uses a model compression method based on knowledge distillation to realize each step.

[0021] By adopting the above technical solution, the application has the following beneficial effects: The model compression method and system based on knowledge distillation provided by the application inherit the soft label of a mature teacher model that has been trained, define a distillation loss function of a student model using the soft label, compress the student model, expand the training set by randomly extracting samples to mix and generate new samples, alleviate the problem of insufficient data, improve the robustness of the model to noise and distribution deviation, dynamically adjust the weight to make the student model quickly fit the knowledge of the teacher model, optimize the student model in the later training period, and avoid overfitting the soft label. Through the test set with a proportion of 30% of unlabeled data, the accuracy of the real label is 78.3%, and the accuracy is improved by 11.2% by using the pseudo target. By using a mature student model for deployment, the intermediate layer feature alignment is omitted, the computing cost is reduced, and the application is suitable for edge devices. The GPU memory occupation is reduced from 3.2GB to 1.8GB, which meets the deployment requirements of edge devices. BRIEF DESCRIPTION OF DRAWINGS

[0022] Figure 1 The flowchart of the method of the application. DETAILED DESCRIPTION

[0023] The technical solutions of the application will be described in detail below with reference to the accompanying drawings. Obviously, the described embodiments are part of the embodiments of the application, rather than all the embodiments. Based on the embodiments in the application, all other embodiments obtained by those skilled in the art without creative labor fall within the scope of the application.

[0024] Figure 1 A specific embodiment of a model compression method based on knowledge distillation of the present application is shown. The case background is that a certain city security system needs to deploy a face recognition model on the edge, and the model is required to run in real time on low-power devices such as AI cameras, while adapting to untrained scenes such as night and side face. The accuracy of the traditional four-layer CNN model in this scene is only 83%, and the model parameters cannot be updated through the cloud. The following steps are included: Step 1: Train the teacher model using the training set until convergence to obtain soft labels corresponding to each training sample.

[0025] In this embodiment, the method for training the teacher model is cross-entropy training, and the soft labels corresponding to each sample are:

[0026]

[0027] wherein, is the i-th probability value of the teacher model output probability, is the i-th probability value of the teacher model output layer original probability value, is the true label index of the i-th probability, and N is the number of teacher model output probabilities.

[0028] In this embodiment, the teacher model is trained using a ResNet-50 pre-trained teacher model, and is trained on ImageNet-10 (10-class subset) until convergence. The soft label loss of each sample is calculated:

[0029] Suppose the true label of a sample is "cat" (index 3), and the teacher model output probability distribution is: [0.1, 0.05, 0.2, 0.6, 0.05] (corresponding to 10 categories), then: .

[0030] Step 2: Initialize the student model and define the distillation loss function of the student model according to the soft label.

[0031] In this embodiment, the student model uses a four-layer CNN (input layer 32x32x3, convolutional layer 32@3x3, fully connected layer 512), and of course other models can also be selected, as long as the teacher model is first trained using the method of the present application, and the mature parameters of the teacher model are used to dynamically train the student model.

[0032] In this embodiment, the distillation loss function of the student model is defined according to the soft label is:

[0033]

[0034]

[0035] wherein, is a soft label supervision loss value, is a divergence loss value, is a temperature parameter for smoothing the probability distribution; is all original probability values of the student model output layer sample; is all original probability values of the teacher model output layer; is a dynamic adjustment weight coefficient, decreases linearly with the training process, so that the student model gradually transitions from relying on teacher knowledge to independent learning, represents the number of times of traversing the complete training data set once, represents the total number of training rounds. By using the dynamic adjustment weight coefficient, the soft label is focused on in the early stage of training , and the real label is gradually focused on in the later stage, and decreases linearly with the training process, so that the student model gradually transitions from relying on teacher knowledge to independent learning, || represents the KL divergence, that is, the difference between two probability distributions.

[0036] When is 50 times, is 100 times, the dynamic weight coefficient is: ; let the temperature T = 4, the teacher model output softmax value is [0.1, 0.4, 0.5], the student model output softmax value is [0.2, 0.3, 0.5], and the divergence loss is: ; then the distillation loss function is: ; Step 3: Randomly extract two samples in the training set, generate a new mixed sample according to a certain mixing ratio, repeat this step, and generate multiple mixed samples.

[0037] In this embodiment, the method of generating a new mixed sample according to a certain mixing ratio is:

[0038]

[0039] wherein, wherein, and two input samples randomly selected from the training data, and respectively, and corresponding true labels, is a mixing ratio, which is subject to a Beta distribution , used to control the mixing intensity. By mixing samples, the training data can be expanded, the problem of insufficient data can be alleviated, and the robustness of the model to noise and distribution shift can be improved.

[0040] Randomly select two training samples: <1> Sample A: input x1 (pedestrian image), label y1 = 0 (pedestrian) <2> Sample B: input x2 (vehicle image), label y2 = 1 (vehicle) Mixing ratio λ = 0.3 (subject to Beta(0.4, 0.4) distribution). Generate mixed sample:

[0041]

[0042] The first pixel value of is [0.69] (normalized to the interval [0, 1]), so .

[0043] The cross-entropy loss generated by the enhancement strategy for the newly generated mixed sample is: =

[0044] = ; where, is the cross-entropy loss function.

[0045] Step 4: Mix the mixed samples and the training set according to the pre-set ratio to enhance the training set. By integrating the mixed samples into the training set, the problem of insufficient training data is solved.

[0046] Step 5: Generate pseudo labels for all mixed samples in the enhanced training set, and all pseudo labels are generated by the teacher model. By having the teacher model generate pseudo labels for unlabeled data, the introduction of pseudo labels enhances data diversity and solves the generalization problem of lightweight models in untrained scenarios.

[0047] In this embodiment, the proportion of unlabeled data is 30%. If unlabeled sample C (a blurred profile image) is processed, the teacher model outputs the probability distribution: [0.15, 0.25, 0.6] (corresponding to the categories: pedestrian, vehicle, cyclist). The pseudo label generated is: , the category is an AUTOBOTS.

[0048] Step 6: Joint training of the student model using the enhanced training dataset to obtain a lightweight student model for deployment.

[0049] In this embodiment, when the student model is jointly trained using the enhanced training set: The enhanced training set sample is input into the student model, the original sample in the training set uses a distillation loss function during training, and the mixed sample uses a cross-entropy loss function for training. The total loss function of joint training is: ; Wherein, is the cross-entropy loss weight.

[0050] Table 1 Performance comparison results of the present application and the traditional method

[0051] On the test set with 30% unlabeled data, the pseudo target improves the accuracy by 11.2% (compared with 78.3% using only real labels). Omitting the intermediate layer feature alignment improves the training speed, and the GPU memory usage is reduced from 3.2GB to 1.8GB. The parameter amount of the present application is significantly reduced, meeting the requirements of lightweight, and meeting the deployment requirements of edge devices. This example verifies that the present application can significantly improve the generalization ability and deployment efficiency in actual scenarios while maintaining the model compression rate.

[0052] The present application also provides a model compression system based on knowledge distillation, which uses a model compression method based on knowledge distillation to realize each step.

[0053] Finally, it should be noted that: the above embodiments are only used to illustrate the technical solutions of the present application, and not to limit them; although the present application has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand: it can still modify the technical solutions recorded in the foregoing embodiments, or make equivalent replacement for part or all of the technical features; and these modifications or replacements do not make the essence of the corresponding technical solutions deviate from the scope of the technical solutions of the embodiments of the present application.

Claims

1. A model compression method based on knowledge distillation, characterized in that: The following steps are involved: Step 1: Use the training set to train the teacher model until convergence, and obtain the soft labels corresponding to each training sample; Step 2: Initialize the student model and define the distillation loss function of the student model based on the soft labels; Step 3: Randomly extract two samples from the training set and generate a new mixed sample according to a certain mixing ratio. Repeat this step to generate multiple mixed samples. Step 4: The mixed samples and the training set are mixed in a pre-set ratio to enhance the training set; Step 5: Generate pseudo labels for all mixed samples in the enhanced training set, where the pseudo labels are generated using the teacher model. Step 6: Use the enhanced training dataset to jointly train the student model and obtain a lightweight student model for deployment.

2. A model compression method based on knowledge distillation according to claim 1, characterized in that: The method for training the teacher model is cross entropy training, and the soft labels corresponding to each sample are: in, is the i-th probability value of the teacher model output probability, is the i-th probability value of the original probability value of the teacher model output layer, is the true label index of the i-th probability, and N is the number of output probabilities of the teacher model.

3. A model compression method based on knowledge distillation according to claim 2, characterized in that: Define the distillation loss function of the student model based on the soft labels for: in, For supervision of soft labels Divergence loss value, is the temperature parameter used to smooth the probability distribution; Output all original probability values ​​of the student model layer samples; All original probability values ​​of the teacher model output layer; To dynamically adjust the weight coefficient, As the training process decreases linearly, the student model gradually transitions from relying on the teacher's knowledge to independent learning. Indicates the number of times the complete training data set is traversed in a single pass, Represents the total number of training rounds.

4. A model compression method based on knowledge distillation according to claim 3, characterized in that: The method of generating a new mixed sample according to a certain mixing ratio is: in, and are two input samples randomly selected from the training data, and The samples and The corresponding true label, is the mixing ratio.

5. A model compression method based on knowledge distillation according to claim 4, characterized in that: The cross entropy loss generated by the enhanced strategy for the newly generated mixed samples is: ； in, is the cross entropy loss function.

6. A model compression method based on knowledge distillation according to claim 5, characterized in that: When jointly training the student model using the augmented training set: Input the enhanced training set samples into the student model, The original samples in the training set are trained using the distillation loss function, and the mixed samples are trained using the cross entropy loss function; The total loss function of joint training is: in, is the cross entropy loss weight.

7. A model compression system based on knowledge distillation, characterized in that: Each step of the model compression method based on knowledge distillation is implemented using any one of claims 1 to 6.

Images