A knowledge distillation method based on adaptive matching of head output features

Abstract

The present invention discloses a knowledge distillation method based on adaptive matching of head features, which sorts and prunes the output layer features of the teacher network and the student network, focuses on the head features obtained after pruning, and calculates the head distillation loss according to the corresponding sorting method; then prunes the feature layer containing the target class information in the head features to obtain the main features, adaptively matches in the main features, and calculates the adaptive output feature loss according to the corresponding sorting method; and uses the total distillation loss to prompt the student network to learn the inter-class similarity information of the main output feature layer. The present invention utilizes the information of non-target class knowledge and integrates the method of adaptive attention matching to locate the main features containing more generalized information, adaptively extracts classification information that is more conducive to the learning of the student network, improves the adaptability of the student network to unfamiliar data sets, and also enhances the generalization and robustness of the model.

Description

Technical Field

[0001] The present invention belongs to computer vision technology, and in particular relates to a knowledge distillation method based on adaptive matching of head output features. Background Art

[0002] Knowledge distillation is to transfer the knowledge of the teacher model to the student model through training. With the development of deep learning, the effect of knowledge distillation has gradually improved. However, since the student model is smaller than the teacher model, the student model has limited expressive power and it is difficult to fully learn more important knowledge from the teacher model. At present, the known knowledge distillation algorithms are mainly based on the following two ideas: one is to transfer the output features of the teacher model as knowledge to the student model; the other is to extract deep features from the middle layer of the teacher network as knowledge distillation to the student model.

[0003] In the article "Decoupled Knowledge Distillation", Zhao Boyui continued the idea of ​​passing the output features of the teacher model as knowledge to the student model, and believed that the target class loss and non-target class loss should be reasonably balanced in the loss function in order to more flexibly and effectively play the role of output layer feature distillation. Therefore, the author adopted a decoupling method to restate the classic distillation loss as the sum of the target class loss and the non-target class loss, solving the problem that the non-target class information is suppressed in the classic distillation loss and it is difficult to fully reflect its value. However, in the distillation process, there is still a problem that due to the difference in the expression ability of the teacher model and the student model, the student model cannot make full use of the information, and cannot extract and match more detailed and important information with limited capabilities. Summary of the invention

[0004] The purpose of the present invention is to provide a knowledge distillation method based on adaptive matching of head output features, which effectively explores some potentials of output layer feature distillation, solves the problems of insufficient utilization of output feature knowledge and large differences in the expressive ability of student models and teacher models, and does not introduce additional parameters to generate computational and memory losses, and has higher training efficiency and better feature portability.

[0005] The technical solution to achieve the purpose of the present invention is: a knowledge distillation method based on adaptive matching of head output features, comprising the following steps:

[0006] Step S1: randomly collect K labeled images from the CIFAR-100 dataset, 10000<K≤60000, normalize the K images, and unify the pixel size to h×w, where h is the image height and w is the image width. The images after unification are randomly divided into a training dataset and a test dataset in a ratio of 5:1, and the training dataset is enhanced to form a teacher-student network training dataset. The teacher network is pre-trained using the teacher-student network training dataset to obtain a pre-trained teacher network, and then the process goes to step S2.

[0007] Step S2: Use the pre-trained teacher network and the student network to extract the original output layer features of the corresponding pre-trained teacher network and the student network respectively, and sort them in order according to the score size to form the corresponding sorted set of the original output layer features of the pre-trained teacher network and the sorted set of the original output layer features of the student network, and then go to step S3.

[0008] Step S3, perform the first pruning processing on the sorted set of original output layer features of the pre-trained teacher network and the sorted set of original output layer features of the student network in step S2, that is, respectively intercept the first m layers of output layer features in their respective sorted sets as the head output features of the pre-trained teacher network and the head output features of the student network, and then proceed to step S4.

[0009] Step S4: Calculate the original distillation loss using the original output layer features of the pre-trained teacher network before sorting and the original output layer features of the student network before sorting in step S2. Then calculate the head distillation loss using the head output features of the pre-trained teacher network and the head output features of the student network obtained in step S3, and proceed to step S5.

[0010] Step S5: Perform a second pruning process on the head output features of the pre-trained teacher network and the head output features of the student network obtained in step S3, that is, decouple the target class head output features and the non-target class head output features and remove the target class head output features. Perform local output feature matching on the retained non-target class head output features, that is, output feature adaptive matching, and calculate the adaptive output feature loss of the non-target class, and then proceed to step S6.

[0011] Step S6: Use the head distillation loss, the original distillation loss, and the adaptive output feature loss of the non-target class to obtain the total distillation loss, and use it to update the network parameters of the student network, finally obtain the trained student network, and go to step S7.

[0012] Step S7: Input the test data set into the trained student network, output the prediction result corresponding to each sample in the test set, and test the accuracy of the trained student network.

[0013] Compared with the prior art, the present invention has the advantages of:

[0014] (1) Compared with the existing knowledge distillation method, only the part with a higher proportion in the probability distribution of the output features is considered, the features that make almost no contribution are cut off, and then the knowledge of non-target classes is refined in the features with an important proportion. The method based on the adaptive matching of head output features described in the present invention solves the problems of insufficient utilization of output feature knowledge and large differences in the expression ability of the student model and the teacher model, and does not introduce additional parameters to generate calculation and memory losses, and has higher training efficiency and better feature portability.

[0015] (2) The present invention proposes to apply the attention method to the knowledge distillation method based on the head features of the output layer, forcing the student network to imitate the more important inter-category similarities of the teacher network, significantly improving the performance of the student network.

[0016] (3) The present invention analyzes and inherits the decoupled knowledge distillation from another perspective, and integrates the method of adaptive attention matching, and finally locates the intermediate subject features containing more generalized information, adaptively extracts the classification information that is more conducive to the student network to learn the teacher network, and also improves the adaptability of the student network to unfamiliar data sets. BRIEF DESCRIPTION OF THE DRAWINGS

[0017] Figure 1 Model diagram of the knowledge distillation method based on adaptive matching of head output features. DETAILED DESCRIPTION

[0018] In order to make the objectives, technical solutions and advantages of the present invention more clear, the embodiments of the present invention are described in further detail below.

[0019] Combination Figure 1 , a knowledge distillation method based on adaptive matching of head output features, comprising the following steps:

[0020] Step S1. Randomly collect K labeled images from the CIFAR-100 data set, 10000<K≤60000, normalize the K images, and unify the pixel size to h×w, where h is the image height and w is the image width; randomly divide the images after unification into a training data set and a test data set in a ratio of 5:1, perform data enhancement on the training data set to form a teacher-student network training data set, use the teacher-student network training data set to pre-train the teacher network, obtain the pre-trained teacher network, and proceed to step S2.

[0021] Step S2: Use the pre-trained teacher network and the student network to extract the original output layer features of the corresponding pre-trained teacher network and the student network respectively, and sort them in order according to the score size to form the corresponding sorted set of the original output layer features of the pre-trained teacher network and the sorted set of the original output layer features of the student network, and then go to step S3.

[0022] Step S3, perform the first pruning process on the sorted set of original output layer features of the pre-trained teacher network and the sorted set of original output layer features of the student network in step S2, that is, intercept the first m layers of output layer features in their respective sorted sets as the head output features of the pre-trained teacher network and the head output features of the student network, as follows:

[0023] First, the original output layer features of the pre-trained teacher network and the original output layer features of the student network are sorted from large to small according to the scores, and the sorted original output layer features of the pre-trained teacher network are sorted into a set O T and the sorted set O of the original output layer features of the student network S Denoted as Where T represents the pre-trained teacher network, S represents the student network, and c is the number of layers in the output layer of the corresponding network.

[0024] Then intercept the sorted set O T The first m output layer features in are used as the head output features of the pre-trained teacher network. Intercept sorted set O S The first m output layer features in are used as the head output features of the student network Defined as Among them, m<c.

[0025] Step S4: Calculate the original distillation loss using the original output layer features of the pre-trained teacher network before sorting in step S2 and the original output layer features of the student network before sorting. Then calculate the head distillation loss using the head output features of the pre-trained teacher network and the head output features of the student network obtained in step S3, as follows:

[0026] Input the original output layer features of the i-th layer of the pre-trained teacher network and the student network into the softmax function layer to obtain the corresponding classification probability p of the i-th class i , that is, p i It represents the probability of the i-th category contained in the teacher-student network training set, which is expressed by the following formula:

[0027]

[0028] In the formula, o iRepresents the output features of the i-th layer. The classification probability set of the pre-trained teacher network The classification probability set of the student network

[0029] The classification probability p corresponding to the original output layer features of the pre-trained teacher network and the original output layer features of the student network is used i , use the KL-Divergence loss function to calculate the classic knowledge distillation loss KD cla , the formula is as follows:

[0030]

[0031] Similarly, the n-th layer head output layer features of the pre-trained teacher network and the student network are input into the softmax function layer to obtain the corresponding n-th class classification probability p n , that is, p n It represents the probability of the nth category included in the teacher-student network training set:

[0032]

[0033] In the formula, o n Represents the output features of the nth layer, n<m. The classification probability set of the pre-trained teacher network The classification probability set of the student network

[0034] Calculate the head distillation loss KD using the head output features of the pre-trained teacher network obtained in step S3 and the head output features of the student network ｈ , the formula is as follows:

[0035]

[0036] There are two ways to sort the head output layer features: sorting by the size of the head output features of the pre-trained teacher network and sorting by the size of the head output layer features of the student network. The size of the head output layer features of the pre-trained teacher network and the size of the head output layer features of the student network after the final sorting under each sorting method are one-to-one corresponding. According to these two sorting methods, two head output feature matching losses are obtained, and these two head output feature matching losses are defined as the probability similarity of sorting according to the size of the head output layer features of the teacher network And the probability similarity sorted by the feature size of the student network head output layer

[0037]

[0038]

[0039] In the formula, is the classification probability of the pre-trained teacher network output corresponding to the sorting of the head output layer features of the pre-trained teacher network, It is the classification probability of the student network output corresponding to the order of the features of the head output layer of the pre-trained teacher network. The classification probability of the pre-trained teacher network output corresponding to the feature sorting of the student network head output layer, The classification probability of the student network output corresponding to the order of the features of the student network head output layer.

[0040] Therefore, the probability similarity loss KD of the head output feature matching H Rebuild as:

[0041]

[0042] The classic knowledge distillation loss KD cla Probabilistic similarity loss KD matching the head output features H Weighted summation is used as the total distillation loss KD of the knowledge distillation method based on adaptive matching of head output features total The influencing factors are as follows:

[0043] KD totaI =α1KD H +α2KD cla

[0044] Where α1 is the probability similarity loss KD for balancing the head output feature matching H The hyperparameters of α2 are the classic knowledge distillation loss KD cla Hyperparameters of .

[0045] Step S5: Perform a second pruning process on the head output features of the pre-trained teacher network and the head output features of the student network obtained in step S3, that is, decouple the target class head output features and the non-target class head output features and remove the target class head output features. Perform local output feature matching on the retained non-target class head output features, that is, output feature adaptive matching, and calculate the adaptive output feature loss of the non-target class, as follows:

[0046] First, the second pruning process is performed on the extracted head output features of the pre-trained teacher network and the student network, that is, the class with the largest classification probability, that is, the target class, is recorded as the tth class. and Respectively expressed as Use the mask to set the output layer features corresponding to the tth class to zero, and retain the output features of the corresponding layer of the non-target class, that is, the main features of the output layer of the pre-trained teacher network are obtained. And the main features of the output layer of the student network

[0047] Then, using the classification probability p n Used as the target prediction value for multi-classification, that is, p n The size indicates the similarity between the nth category and the target category tth category in the image.

[0048] Finally, another round of pruning is performed when calculating the loss. That is, the last layer of features is removed after each round of loss calculation, and the proportion of the remaining output layers with larger non-target class probability distribution in the total number of non-target class output layers is calculated. Finally, an adaptive method is used to retain the output features of the last x layers.

[0049] Probability similarity sorted by the order of output feature size of the last x layers of the pre-trained teacher network and the probability similarity in the order of the output feature size of the last x layers of the student network The formula is as follows:

[0050]

[0051]

[0052] In the formula, is the classification probability of the pre-trained teacher network output corresponding to the feature sorting of the last x layers of the pre-trained teacher network, is the classification probability of the student network output corresponding to the feature sorting of the last x layers of the pre-trained teacher network. The classification probability of the pre-trained teacher network output corresponding to the feature sorting of the last x layers of the student network, The classification probability of the student network output corresponding to the order of the output features of the last x layers of the student network.

[0053] The probability similarity loss KD of the adaptive output feature loss ad Rebuild as:

[0054]

[0055] Step S6: Use the head distillation loss, the original distillation loss, and the adaptive output feature loss of the non-target class to obtain the total distillation loss, and use it to update the network parameters of the student network to finally obtain a trained student network. The total distillation loss calculation process is as follows:

[0056] KD total =α1KD H +α2KD cla +α3KD ad

[0057] Among them, α1 is the probability similarity loss KD of balancing head output feature matching H The hyperparameters of α2 are the classic knowledge distillation loss KD cla The hyperparameters of α3 are the knowledge distillation loss KD based on adaptive matching of head output features. total In order to balance the probability similarity loss KD of the adaptive output feature loss ad Hyperparameters introduced.

[0058] Step S7: input the test set into the trained student network, output the prediction result corresponding to each sample in the test set, and test the accuracy of the trained student network.

[0059] Example 1

[0060] The distillation method based on adaptive matching of head output features described in the present invention comprises the following steps:

[0061] Step S1, randomly collect 60,000 labeled images in the CIFAR-100 data set, normalize the 60,000 images, unify the pixel size to 32×32, randomly divide the images after the unified size into a training data set and a test data set at a ratio of 5:1, perform data enhancement on the training data set to form a teacher-student network training data set, use the teacher-student network training data set to pre-train the teacher network to obtain a pre-trained teacher network, wherein the data enhancement operation includes image scaling and random flipping, the image scaling ratio is inward and outward scaling according to 10% of the original image, the random flipping angle is between -20° and 20°, and the number of image categories is 100.

[0062] Step S2: Use the pre-trained teacher network and the original student network to extract the original output layer features of the corresponding pre-trained teacher network and the original student network, and sort them in order according to the score size to form a sorted set of the corresponding pre-trained teacher network original output layer features and a sorted set of the original student network original output layer features.

[0063] Step S3, perform the first pruning processing on the sorted set consisting of the original output layer features of the pre-trained teacher network in step S2 and the sorted set consisting of the original output layer features of the student network, that is, respectively extract the first 45 layers of output layer features in their respective sorted sets as the head output features of the pre-trained teacher network and the head output features of the student network.

[0064] Step S4: Calculate the original distillation loss using the original output layer features of the pre-trained teacher network and the original output layer features of the student network before sorting in step 2. Calculate the head distillation loss using the head output features of the pre-trained teacher network and the head output features of the student network in step 3.

[0065] Step S5: Perform a second pruning process on the head output features of the pre-trained teacher network and the head output features of the student network in step 3, that is, decouple the target class head output features and the non-target class head output features and remove the target class head output features. Perform local output feature matching on the retained non-target class head output features, that is, output feature adaptive matching, and calculate the adaptive output feature loss of the non-target class.

[0066] When calculating the loss, another round of pruning is performed, that is, each round of loss calculation removes the last layer of features, calculates the proportion of the remaining output layers with larger non-target probability distribution in the total number of non-target output layers, and finally uses an adaptive method to retain the output features of the last 5 layers

[0067] Step S6: Use the head distillation loss, the original distillation loss, and the adaptive output feature loss of the non-target class to obtain the total distillation loss, and use it to update the network parameters of the student network, finally obtain the trained student network, and go to step S7.

[0068] Step S7: input the student network test set into the trained student network, output the prediction result corresponding to each sample in the test set, and test the accuracy of the trained student network.

[0069] The method of the present invention uses python programming language and pytorch framework language to build a network framework on an Nvidia 2080Ti GPU host for relevant experiments. For classification tasks, we calculate the sum of the three losses from the original knowledge distillation loss, the probability similarity loss of the head output feature, and the probability similarity loss of the adaptive output feature loss. The method of the present invention uses three hyperparameters α1, α2, and α3 to balance the distillation loss of the equation. Set the hyperparameters {α1=0.5, α2=0.5, α3=0.1} and the distillation temperature T=4 for classification experiments. We use the SGD optimizer to train all models for 240 epochs, with a momentum of 0.9 and a weight decay of 0.0001. We initialize the learning rate to 0.05 and decay it every 30 epochs. Multiple training on the training sample set resulted in a distillation network model based on adaptive matching of head output features.

Claims

1. A knowledge distillation method based on adaptive matching of head output features, characterized in that: The steps are as follows: Step S1: Randomly collect K labeled images in the CIFAR-100 dataset, where 10000 < K ≤ 60000. Normalize the above K images, unify the pixel size to h×w, where h is the image height and w is the image width; randomly divide the images with unified size into a training dataset and a test dataset according to a ratio of 5:

1. Perform data augmentation on the training dataset to form a teacher-student network training dataset, and use the teacher-student network training dataset to pre-train the teacher network to obtain a pre-trained teacher network, then proceed to step S2; Step S2: Use the pre-trained teacher network and the student network to extract the original output layer features of the corresponding pre-trained teacher network and student network respectively, and sort them in descending order of scores respectively to form the sorted sets of the original output layer features of the pre-trained teacher network and the sorted sets of the original output layer features of the student network, then proceed to step S3; Step S3: Perform the first pruning process on the sorted set of the original output layer features of the pre-trained teacher network and the sorted set of the original output layer features of the student network in step S2, that is, intercept the first m layers of output layer features in their respective sorted sets as the head output features of the pre-trained teacher network and the head output features of the student network, then proceed to step S4; Step S4: Calculate the original distillation loss using the original output layer features of the pre-trained teacher network before sorting and the original output layer features of the student network before sorting in step S2; then calculate the head distillation loss using the head output features of the pre-trained teacher network and the head output features of the student network obtained in step S3, then proceed to step S5; Step S5: Perform the second pruning process on the head output features of the pre-trained teacher network and the head output features of the student network obtained in step S3, that is, decouple the target class head output features and non-target class head output features and remove the target class head output features; perform local output feature matching on the remaining non-target class head output features, that is, adaptive output feature matching, and calculate the adaptive output feature loss of the non-target class, then proceed to step S6; Step S6: Use the head distillation loss, the original distillation loss, and the adaptive output feature loss of the non-target class to obtain the total distillation loss, and use this to update the network parameters of the student network, and finally obtain a trained student network, then proceed to step S7; Step S7: Input the test dataset into the trained student network, output the prediction results corresponding to each sample in the test set, and test the accuracy of the trained student network.

2. The knowledge distillation method based on adaptive matching of head output features according to claim 1, characterized in that: In step S3, perform the first pruning process on the sorted set of the original output layer features of the pre-trained teacher network and the sorted set of the original output layer features of the student network in step S2, that is, intercept the first m layers of output layer features in their respective sorted sets as the head output features of the pre-trained teacher network and the head output features of the student network, specifically as follows: First, the original output layer features of the pre-trained teacher network and the original output layer features of the student network are sorted from large to small according to the scores, and the sorted original output layer features of the pre-trained teacher network are sorted into a set O T and the sorted set O of the original output layer features of the student network S Denoted as Where T represents the pre-trained teacher network, S represents the student network, and c is the number of layers in the output layer of the corresponding network; Then intercept the sorted set O T The first m output layer features in are used as the head output features of the pre-trained teacher network. Intercept sorted set O S The first m output layer features in are used as the head output features of the student network Defined as Where m <c。 3. The knowledge distillation method based on adaptive matching of head output features according to claim 2 is characterized in that: In step S4, the original output layer features of the pre-trained teacher network and the original output layer features of the student network before the sorting in step S2 are used to calculate the original distillation loss; and the head output features of the pre-trained teacher network and the head output features of the student network obtained in step S3 are used to calculate the head distillation loss, as follows: Input the original output layer features of the i-th layer of the pre-trained teacher network and the student network into the softmax function layer to obtain the corresponding classification probability p of the i-th class i , that is, p i It represents the probability of the i-th category contained in the teacher-student network training set, which is expressed by the following formula: In the formula, o i Represents the output features of the i-th layer; the classification probability set of the pre-trained teacher network The classification probability set of the student network The classification probability p corresponding to the original output layer features of the pre-trained teacher network and the original output layer features of the student network is used i , use the KL-Divergence loss function to calculate the classic knowledge distillation loss KD cla , the formula is as follows: Similarly, the n-th layer head output layer features of the pre-trained teacher network and the student network are input into the softmax function layer to obtain the corresponding n-th class classification probability p n , that is, p n It represents the probability of the nth category included in the teacher-student network training set: where, o n represents the output feature of the n-th layer, n < m; the set of classification probabilities of the pre-trained teacher network the set of classification probabilities of the student network Calculate the head distillation loss KD using the head output features of the pre-trained teacher network obtained in step S3 and the head output features of the student network H , the formula is as follows: There are two ways to sort the head output layer features: sorting by the size of the head output features of the pre-trained teacher network and sorting by the size of the head output layer features of the student network; the size of the head output layer features of the pre-trained teacher network and the size of the head output layer features of the student network after the final sorting under each sorting method are one-to-one corresponding; according to these two sorting methods, two head output feature matching losses are obtained, and these two head output feature matching losses are defined as the probability similarity of sorting according to the size of the head output layer features of the teacher network And the probability similarity sorted by the feature size of the student network head output layer In the formula, is the classification probability of the pre-trained teacher network output corresponding to the sorting of the head output layer features of the pre-trained teacher network, The classification probability of the student network output corresponding to the order of the features of the head output layer of the pre-trained teacher network; The classification probability of the pre-trained teacher network output corresponding to the feature sorting of the student network head output layer, The classification probability of the student network output corresponding to the feature sorting of the student network head output layer; Therefore, the probability similarity loss KD of the head output feature matching H Rebuild as: The classic knowledge distillation loss KD cla Probabilistic similarity loss KD matching the head output features H Weighted summation is used as the total distillation loss KD of the knowledge distillation method based on adaptive matching of head output features total The influencing factors are as follows: KD total =α1KD H +α2KD cla Where α1 is the probability similarity loss KD for balancing the head output feature matching H The hyperparameters of α2 are the classic knowledge distillation loss KD cla Hyperparameters of .

4. The knowledge distillation method based on adaptive matching of head output features according to claim 3 is characterized in that: In step S5, the head output features of the pre-trained teacher network and the head output features of the student network obtained in step S3 are pruned for the second time, that is, the target class head output features and the non-target class head output features are decoupled and the target class head output features are removed; the retained non-target class head output features are matched with local output features, that is, output feature adaptive matching, and the non-target class knowledge distillation loss is calculated, as follows: First, the second pruning process is performed on the extracted head output features of the pre-trained teacher network and the student network, that is, the class with the largest classification probability, that is, the target class, is recorded as the tth class. and Respectively expressed as Use the mask to set the output layer features corresponding to the tth class to zero, and retain the output features of the corresponding layer of the non-target class, that is, the main features of the output layer of the pre-trained teacher network are obtained. And the main features of the output layer of the student network Then, using the classification probability p n Used as the target prediction value for multiple classifications, that is, p n The size indicates the similarity between the nth category and the target category tth category in the image; Finally, another round of pruning is performed when calculating the loss. That is, the last layer of features is removed after each round of loss calculation, and the proportion of the remaining output layers with larger non-target probability distribution in the total number of non-target output layers is calculated. Finally, the output features of the last x layers are retained using an adaptive method. Probability similarity sorted by the order of output feature size of the last x layers of the pre-trained teacher network and the probability similarity in the order of the output feature size of the last x layers of the student network The formula is as follows: In the formula, is the classification probability of the pre-trained teacher network output corresponding to the feature sorting of the last x layers of the pre-trained teacher network, is the classification probability of the student network output corresponding to the feature sorting of the last x layers of the pre-trained teacher network; The classification probability of the pre-trained teacher network output corresponding to the feature sorting of the last x layers of the student network, The classification probability of the student network output corresponding to the order of the output features of the last x layers of the student network; The probability similarity loss KD of the adaptive output feature loss ad Rebuild as:

5. The knowledge distillation method based on adaptive matching of head output features according to claim 4 is characterized in that: In step S6, the head distillation loss, the original distillation loss, and the adaptive output feature loss of the non-target class are used to obtain the total distillation loss, that is, the knowledge distillation loss KD based on the adaptive matching of the head output features. total The formula is: KD total =α1KD H +α2KD cla +α3KD ad Among them, α1 is the probability similarity loss KD of balancing head output feature matching H The hyperparameters of α2 are the classic knowledge distillation loss KD cla The hyperparameters of α3 are the knowledge distillation loss KD based on adaptive matching of head output features. total In order to balance the probability similarity loss KD of the adaptive output feature loss ad Hyperparameters introduced.

Images