An image recognition model compression method based on adversarial distillation technology

Abstract

The present invention discloses an image recognition model compression method based on adversarial distillation technology, which belongs to the field of computer vision. The present invention introduces the adversarial learning idea into knowledge distillation, and adopts adversarial generative network (GAN) as a framework. The generator is mapped to a student model, and the real sample is mapped to the knowledge of the teacher model to perform cyclic adversarial training. Specifically, it includes: collecting image recognition data sets and labeling categories; using GAN to generate image data sets; selecting teacher models and training; using optimized or condensed small network versions to select student models; selecting knowledge types; introducing adversarial distillation algorithms to train student models; realizing image recognition tasks, and solving the problems of large network model parameters and low computational efficiency.

Description

Technical Field

[0001] The invention relates to an image recognition model compression method based on knowledge distillation technology, and belongs to the field of computer vision. Background Art

[0002] Image recognition technology refers to the technology of using computers to process, analyze and understand images in order to identify targets and objects of various different patterns.

[0003] Deep learning technology is a new research direction in the field of machine learning and has achieved great success in industry and academia, mainly due to its scalability in encoding large-scale data and manipulating billions of model parameters.

[0004] Model compression technology is a key technology for effectively deploying neural network models on mobile devices with limited computing resources and tight energy budgets.

[0005] Currently, many studies have been conducted to improve the computational efficiency of deep neural networks through model compression. Deep model compression methods include parameter pruning and sharing, low-rank factorization, transfer or compact convolutional filters, and knowledge distillation.

[0006] Knowledge distillation is the most representative model compression, which can effectively transfer the knowledge in a large teacher model to a small student model. Knowledge transfer usually transfers the potential information contained in the teacher model to the student model, which provides technical support for embedding deep learning-based computer vision technology into small devices.

[0007] The knowledge distillation system is mainly composed of three main parts: knowledge, distillation algorithm, and teacher-student architecture. Knowledge is the mapping learned by the model from input vector to output vector. The distillation algorithm is the process of knowledge transfer. The teacher-student architecture is a universal carrier for knowledge transfer.

[0008] The "temperature" mechanism is to add a temperature factor T to the Softmax layer in the final network output layer to control the importance of information in its output results. This allows the student model to learn not only the distribution information of the label data encoded with one-hot, but also the potential information brought by the probability of each label in the teacher model.

[0009] Generative Adversarial Networks (GANs) are one of the most promising methods for unsupervised learning on complex distributions. In adversarial learning, the discriminator estimates the distribution probability from the training data, while the generator attempts to deceive the discriminator with generated data samples. Adversarial learning can generate data through the generator, which plays a role in data enhancement.

[0010] In knowledge distillation, it is difficult for the teacher model to fully learn the probability distribution of data from real data. In addition, even if different combinations of knowledge types are selected, the student model cannot accurately imitate the teacher model due to its small capacity.

[0011] The adversarial distillation algorithm, by drawing on the idea of ​​adversarial learning, not only solves the problem that the student model cannot completely imitate the teacher model due to its small model parameters, but also the performance of the student model under adversarial distillation is significantly improved compared to simply training the student model, allowing the teacher model and the student model to better learn the real data distribution. Summary of the invention

[0012] In knowledge distillation based on feature or response knowledge, the small capacity of the student model network makes it unable to fully imitate the teacher network. The present invention adopts a knowledge distillation algorithm based on adversarial distillation, which can ensure that the teacher model and the student model better understand the distribution of real data and solve the problem that the student model network cannot fully imitate the teacher model network due to its small capacity.

[0013] Adversarial distillation introduces the idea of ​​adversarial learning, using adversarial generative networks as the model framework, treating the student model as the generator, and the teacher model providing real samples and introducing the "teaching assistant" mechanism as the discriminator. The type of knowledge in adversarial distillation adopts feature-based knowledge, setting the transferred knowledge as the feature map of the middle layer of the network and the output result of the classification layer. The generator generates knowledge information according to the setting of the knowledge type, and uses the "teaching assistant" system to determine whether the information generated by the generator is the information generated by the teacher model. This allows the student network to fully learn the data distribution of the knowledge generated in the teacher network. In addition, the accuracy of the student model in adversarial distillation will be further improved compared to the accuracy of the independently trained student model, and the model can be optimized for loss faster.

[0014] The purpose of the present invention is to combine deep learning with knowledge distillation, and use adversarial distillation technology to compress deep learning models for image recognition tasks. The deep learning model has too many calculation parameters, which makes the calculation speed too low. In the industrial field, it is a great challenge to deploy it on a device with limited resources. By compressing the model with the help of knowledge distillation technology, the problem of large model calculation parameters and slow calculation speed can be solved within an acceptable range with limited loss accuracy.

[0015] The specific technical solutions are as follows:

[0016] A method for compressing an image recognition model based on adversarial distillation technology, characterized by comprising the following steps:

[0017] 1. According to different image recognition tasks, data collection, sample labeling and definition of evaluation indicators for image recognition tasks are carried out.

[0018] 2. Use the generative adversarial network to generate data for the collected data set and enhance the generalization of data distribution.

[0019] 3. The selection of teacher-student structure is divided into three parts:

[0020] (1) Selection of teacher model

[0021] The teacher model network used in the present invention is ResNet, in which the residual module and the batch normalization module are introduced to avoid model degradation and gradient disappearance and explosion during the deeper construction process of the network, thereby achieving higher image recognition accuracy.

[0022] (2) Selection of knowledge types

[0023] In the model network, the feature graph in the middle layer and the results of the Softmax layer in the classification layer contain rich information. The feature graph in the middle layer is used to directly match the feature activation of the teacher model and the student model. The results of the Softmax layer are used to generate "dark knowledge" for the student model, so that the student model learns the probability distribution of each data label, which plays a certain role in label smoothing. The knowledge type in the present invention selects the knowledge type based on the feature.

[0024] (3) Selection of student models

[0025] According to the network structure of ResNet in the teacher model and the selection of knowledge types, the VGG-16 model is used as the student model. First, VGG-16 and ResNet can effectively match the feature map size, and there is no need to use a conversion function to convert the feature maps of the two to a uniform size to prevent the loss of feature map information. Secondly, the convolution + convolution + convolution + pooling module is used to replace the traditional convolution module to solve the problem of large parameters in the traditional convolution module and prevent overfitting.

[0026] 4. Adversarial Distillation Model

[0027] The adversarial distillation model uses the student model as the generator, and the knowledge transferred from the teacher model is used to represent the real data. The "teaching assistant" mechanism is used as the discriminator. The model mainly includes two steps:

[0028] (1) Define the adversarial distillation model structure

[0029] Generator

[0030] The generator uses the VGG-16 network, denoted as G. According to the set knowledge type, the generator will generate feature maps based on the intermediate layer and the output results of the Softmax layer in the classification layer, denoted as G1 and G2 respectively.

[0031] Discriminator

[0032] The discriminator uses the LeNet convolutional structure, denoted as D, which concatenates the feature map of the middle layer of the student model network and the feature map of the teacher model network as input to judge the authenticity of the data.

[0033] Real samples

[0034] The feature maps of the intermediate network layers of the teacher model are used as knowledge to construct real samples.

[0035] (2) Define the training process and end point

[0036] The generator and discriminator are updated alternately and cyclically, one is fixed and the other is updated, and a "temperature" mechanism is introduced to mine the potential information of the output results of the Softmax layer in the classification layer. After several cycles of training, the probability that the discriminator judges that the generated data is real data is 50% as the training end point.

[0037] 5. Model training and testing

[0038] (1) Use adversarial distillation to train the student model network and adjust its parameters.

[0039] (2) After training, the student network is used to perform image recognition verification on the test set, and its performance is evaluated using evaluation indicators. BRIEF DESCRIPTION OF THE DRAWINGS

[0040] In order to make the solution of the present invention clearer, the invention will be further described below with reference to the accompanying drawings.

[0041] Figure 1 This is the model compression structure based on the adversarial distillation algorithm of the present invention.

[0042] Figure 2 This is a flow chart of the adversarial distillation algorithm of the present invention.

[0043] Figure 3 This is a knowledge distillation structure diagram based on response knowledge in the present invention. DETAILED DESCRIPTION

[0044] To achieve the above purpose, the technical solution adopted by the present invention is an image recognition model compression method based on adversarial distillation technology. The method specifically comprises the following steps:

[0045] First, according to the task requirements, the experimental image dataset required by the task is collected and the images are annotated. Among them, the original image dataset is input into the adversarial generative network to generate the image dataset. Secondly, the teacher-student framework is selected, and the process is divided into: selecting the teacher model and training it, selecting a student model using the optimized\condensed small network version, and selecting the knowledge type. Finally, the adversarial distillation algorithm is introduced to train the student model.

[0046] 1. Dataset

[0047] (1) Datasets and evaluation metrics

[0048] Task: Mobile phone brand image recognition task. The dataset is a manually collected mobile phone brand dataset with a size of 5742 pictures. The labelers labeled the pictures with the labels Huawei, Apple, OPPO, vivo and Xiaomi. The specific data statistics are shown in Table 1 below.

[0049] Table 1 Collection of mobile phone brand datasets

[0050]

[0051] The image recognition task uses two evaluation indicators to evaluate the performance of the model, Accuracy (Acc) and F1-score. ACC is used to measure the percentage of images predicted correctly by the model in the total image data. The F1 value is defined as the harmonic mean of precision and recall.

[0052]

[0053]

[0054]

[0055]

[0056] P represents the number of positive samples. N represents the number of negative samples. TP represents the number of samples predicted by the model to be positive but are actually positive samples. TN represents the number of samples predicted by the model to be negative but are actually negative samples. FP represents the number of samples predicted by the model to be positive but are actually negative samples. FN represents the number of samples predicted by the model to be negative but are actually positive samples.

[0057] 2. GAN generates image datasets

[0058] GAN is used to generate artificial synthetic data sets to further enhance the generalization of data. The steps are as follows:

[0059] Generator: By inputting a random noise z, adding the noise z to the generator, and outputting an image, denoted as G(z).

[0060] Discriminator: Determines whether a piece of image data is real. The input is a picture x, and the output is the probability that the picture is a real picture, denoted as D(x).

[0061] First, set the number of alternating updates T = 10. Train the generator and fix the discriminator to update the parameters of the generator. After that, train the discriminator and fix the generator to update the parameters of the discriminator. Continue to alternately update the parameters of the generator and discriminator until the last time the discriminator has an accuracy of 50% for its input image, maximizing the benefits of both networks and ending the training. The specific objective function is as follows:

[0062]

[0063] The specific generated data is shown in Table 2 below:

[0064] Table 2 GAN generated dataset

[0065]

[0066] (2) Selection of teacher-student structure

[0067] The selection of the teacher-student framework is divided into three parts: the selection of the teacher model, the selection of the student model, and the selection of the type of knowledge.

[0068] 1. Selection of teacher model

[0069] The teacher model uses the ResNet network architecture. The ResNet network has two major advantages: residual structure and Batch Normalization to solve the degradation problem and gradient vanishing or gradient exploding problem in traditional neural networks after being stacked to a certain depth, so that the model can learn the characteristics of real data as much as possible. The residual structure in ResNet is: the feature map output in the convolutional network layer is recorded as X, firstly, it is reduced in dimension by a 1*1 convolution kernel, then feature extraction is performed by a 3*3 convolution layer, and then the dimension of the initial channels is restored by a 1*1 convolution kernel. Finally, the original feature map X and the output after convolution are added and input into the relu activation function.

[0070] 2. Selection of knowledge type

[0071] Using feature-based knowledge as the knowledge carrier, the feature map output by the middle layer of the teacher model and the output value of the Softmax layer are transferred to the student model as knowledge. Its loss function is:

[0072]

[0073] where f t (x), f s (x) represents the feature map of the intermediate layer between the teacher model and the student model. Φ t (f t (x)), Φ s (f s (x) indicates that when the feature maps of the intermediate layers of the teacher model and the student model have different sizes, they are made to the same size through a conversion function. Represents the similarity function between the feature maps of the student model and the teacher model, including cross entropy loss, maximum mean difference loss, and L1 and L2 norms.

[0074] 3. Selection of student models

[0075] Although the ResNet selected as the teacher model has great advantages in performance, the number of computational parameters of the model is too large, resulting in a slow calculation speed. According to the output scale of the residual structure in each layer of ResNet, which is 112, 56, 28, 14, and 7, it is most appropriate to select the VGG-16 network architecture as the student model. First, the output scale of VGG-16 is consistent with the output scale of ResNet-34, and any feature-based knowledge can be learned. Secondly, VGG-16 is different from the traditional convolution + pooling operation, and uses the convolution + convolution + convolution + pooling operation. This can not only reduce the number of parameters and speed up the calculation, but also prevent overfitting.

[0076] (3) Adversarial Distillation Algorithm

[0077] The adversarial distillation algorithm uses the GAN architecture. The generator is the student model, denoted by G. The discriminator uses the "teaching assistant" network, denoted by D. The transferred knowledge in the teacher network is used as a real sample, denoted by X. The student model is the generator, which is used to generate the knowledge in the corresponding teacher model network as the generated sample. Depending on the type of knowledge, two generators G1 and G2 are set. The G1 generator is used to generate features in the middle layer of the network, and G2 is used to generate the results output by the network Softmax layer. The discriminator D is used to determine the probability of whether the sample input to G1 is a real sample. The specific steps are as follows:

[0078] 1. Input: Given a teacher model network N tAnd its training data set X, the label Y corresponding to the training set and the number of alternating updates Epoch. In addition, a temperature mechanism is introduced, and its variable is denoted by τ.

[0079] 2. Initialize a GAN where the generator (student model) G includes G1, G2, and the discriminator (teacher model) is D.

[0080] 3.Repeat:

[0081] 4. Randomly select an image sample x i and its label y i .

[0082] 5. Substitute the random image samples into the teacher model network to obtain the intermediate layer feature map z T And the output o of the Softmax layer in the output layer T .

[0083] 6. Substitute the random image sample into the generator G1 and G2 to output the intermediate feature map z S And the output o of the Softmax layer in the output layer S .

[0084] 7. Generate the intermediate feature map z from G1 in the generator S Substitute the generated sample into the discriminator D and calculate the probability of whether it is a real sample. It is recorded as D(G1(x i )). At the same time, update the parameters of the discriminator D.

[0085] 8. Introduce the temperature mechanism. In the teacher model and the student model, the temperature mechanism is introduced in the Softmax layer in the output layer, and its output result is used as the "dark knowledge" in the teacher-student architecture. The calculation formula is as follows:

[0086]

[0087]

[0088] Among them, θ S ,θ T is the classification parameter in the student model and teacher model network, and the feature map z S , z T After multiple full connections, it is converted into the numerical value corresponding to each label.

[0089] 9. Calculate D(G1(x i )) and τ(o s ), τ(o T ) to calculate L GAN The loss function is:

[0090]

[0091] Where H(·,·) is used to ensure the cross entropy loss of the student model network performance, λ is a weight parameter that balances the cross entropy loss of the output results in the teacher and student models and the cross entropy loss of the student network output results and labels. γ is a weight parameter to balance the difference between classification loss and adversarial loss.

[0092] 10. Use gradient descent to update the parameters of G1 and G2.

[0093] 11. Update the parameters of the generator and discriminator alternately until the last time the discriminator D determines that the probability of the generator G being real data is close to 50%, and the iteration ends.

[0094] (4) Model training and testing

[0095] In model training, datasets for given tasks and datasets generated by GAN are used for image enhancement.

[0096] Generator parameter update. Calculate the probability of each time the discriminator determines whether the information generated by the generator (student model network) is real data. When the probability is 50%, it means that the discriminator cannot determine whether the generated information is real data, and save the model.

[0097] The discriminator parameters are updated. The model is saved until the discriminator can distinguish between real and generated data sets, that is, the probability of judging that the generated information is real data is less than 20%.

[0098] The discriminator and generator are updated alternately, and the student model updated last time is used as the final image recognition model.

[0099] During the model training process, the channels of each layer of the teacher model network ResNet-34 and the student model network VGG-16 are 112, 56, 28, 14, and 7. λ is 2, the temperature τ is 0.5, and γ is 0.15. The optimization algorithm used in the model training process is stochastic gradient descent, with a learning rate of 0.01. The deep learning framework used is Pytorch, and the machine configuration is NAVIDA3090Ti for model training.

Claims

1. A method for compressing an image recognition model based on knowledge distillation technology, characterized in that: The following steps are involved: (1) Data collection and enhancement Collect pictures according to the given image recognition task and manually annotate their labels; Use the Generative Adversarial Network (GAN) to create synthetic data sets and enhance the generalization of data. The steps include two parts, namely, training the generator to artificially generate images and fixing the discriminator to update the parameters of the generator; then, training the discriminator to determine the probability of whether it is a real image and fixing the generator to update the parameters of the discriminator; The parameters of the generator and discriminator are updated alternately until the benefits of both networks are maximized, and the training ends; (2) Selection of teacher-student structure Selection of teacher model; the teacher model network used is ResNet; residual module and batch normalization module are introduced to enable the network to be constructed at a deeper level; The feature maps in the middle layer are used to directly match the feature activations of the teacher model and the student model; the results of the Softmax layer are used by the student model to generate "dark knowledge"; the knowledge type selection is based on the feature type; Selection of student model: Based on the ResNet network structure and knowledge type in the teacher model, the VGG-16 model is the most appropriate; VGG-16 and ResNet match the feature map size, and use the conversion function to convert the feature maps of the two to a uniform size to prevent the loss of feature map information; VGG-16 uses convolution + convolution + convolution + pooling module to replace the convolution module; (3) Adversarial Distillation Algorithm The adversarial distillation algorithm uses the GAN architecture. The generator is the student model, denoted by G; the discriminator uses the "teaching assistant" network, denoted by D; the transferred knowledge in the teacher network is used as the real sample, denoted by x; the student model is the generator, which is used to generate the knowledge in the corresponding teacher model network as the generated sample; Two generators G1 and G2 are set according to the type of knowledge; the G1 generator is used to generate the features in the middle layer of the network, and G2 is used to generate the results output by the network Softmax layer; the discriminator D is used to determine the probability of whether the sample input by G1 is a real sample; The specific steps are as follows: Given a teacher model network N t And its training data set X and the label Y corresponding to the training set. In addition, the temperature mechanism is introduced, and its variable is denoted as τ; Initialize a GAN, where the generator or student model G includes G1, G2, and the discriminator or teacher model is D; Randomly select an image sample x i and its label y i ; Substitute the random image sample into the teacher model network to obtain the intermediate layer feature map z T And the output o of the Softmax layer in the output layer T ; Substitute the random image sample into the generator G1, G2 to output the intermediate feature map z S And the output o of the Softmax layer in the output layer S ; Generate the intermediate feature map z from the generator G1 S Substitute it into the discriminator D as the generated sample and calculate the probability of whether it is a real sample; record it as D(G1(x i )); while updating the parameters of the discriminator D; Introducing the temperature mechanism; In the teacher model and the student model, the temperature mechanism is introduced in the Softmax layer in the output layer, and the soft targets output by it serve as the "dark knowledge" in the teacher-student architecture; Use gradient descent to update the parameters of G1 and G2; The generator and discriminator parameters are updated alternately until the last time the discriminator D determines that the generator G is the true data, the probability is close to 50%, and the iteration ends; (4) Model training and testing Use adversarial distillation to train the student model network and adjust parameters; After training, the student model network is used to verify image recognition on the test set, and evaluation indicators are used to evaluate its performance.

Images