A model compression method based on attention-like transmission

Abstract

The application discloses a model compression method based on attention-like transmission, comprising: obtaining a to-be-recognized image; using a trained student model to recognize a target object in the to-be-recognized image to obtain a category of the target object; the trained student model is obtained by training a to-be-trained student model according to a prediction score output by the to-be-trained student model based on a training sample, a real label corresponding to the training sample, a first attention-like map output by the to-be-trained student model itself, and a second attention-like map output by a pre-trained teacher model; the to-be-trained student model is trained to obtain the trained student model; a full connection layer connected after a global pooling layer in the to-be-trained student model and the teacher model is replaced by a one-dimensional convolution kernel connected before the global pooling layer. Compared with other model compression methods, the application has higher interpretability, and the application can improve the classification accuracy of the model.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of machine vision, and particularly relates to a model compression method based on class attention transmission. BACKGROUND

[0002] The performance of a neural network model is affected by the model parameter capacity, and a model with more parameters can often obtain better performance. However, a large model has the disadvantages of high hardware requirements and slow running speed, which makes it difficult to deploy to small and mobile devices. The knowledge distillation technology aims to improve the performance of a small model, and its implementation is to additionally let the small model imitate and learn the intermediate output of a pre-trained large model with a larger parameter capacity during the training of the small model. The trained small model is called a student model, the pre-trained large model is called a teacher model, and the intermediate output provided by the large model is called a soft label. Existing knowledge distillation methods can be divided into three categories according to the production method of the soft label they transmit: transmitting prediction scores, transmitting feature maps of the model, and transmitting attention maps of the model.

[0003] Among the existing methods, the knowledge distillation methods based on the transmission of prediction scores and feature maps have good effects. However, since the prediction score values of the neural network model and the feature maps of the intermediate output still cannot be well explained, the knowledge distillation methods based on the transmission of prediction scores and feature maps have poor interpretability. In comparison, the method based on the transmission of attention has good interpretability. However, compared with the methods based on the transmission of prediction scores and feature maps, the existing knowledge distillation methods based on the transmission of attention have poor effects. SUMMARY

[0004] In order to solve the above problems in the related art, the application provides a model compression method based on class attention transmission. The technical problems to be solved by the application are solved by the following technical solutions:

[0005] The application provides a model compression method based on class attention transmission, comprising:

[0006] obtaining a to-be-recognized image;

[0007] using the trained student model to recognize a target object in the to-be-recognized image to obtain the class of the target object;

[0008] The trained student model is obtained by training the student model to be trained according to a prediction score output by the student model to be trained, a real label corresponding to a training sample, a first type of attention map output by the student model to be trained itself, and a second type of attention map output by a pre-trained teacher model; and a fully connected layer connected after a global pooling layer in the student model to be trained and the teacher model is replaced by a one-dimensional convolution kernel connected before the global pooling layer.

[0009] In some embodiments, the student model and the teacher model are both models for target classification.

[0010] In some embodiments, the student model to be trained and the teacher model both include a convolution layer, a one-dimensional convolution kernel, and a global pooling layer; wherein the convolution layer is used for feature extraction of an input image to obtain a feature map of the input image; the one-dimensional convolution kernel is used for generating a class attention map according to the feature map; and the global pooling layer is used for converting the class attention map into a prediction score.

[0011] In some embodiments, before the trained student model is used to identify the target object in the image to be identified, the method further includes:

[0012] Obtaining a plurality of sample images with real labels; the real labels represent the real classes of objects in the sample images;

[0013] Obtaining a student model to be trained for the t-th time and the pre-trained teacher model; wherein when t is 1, the student model to be trained is an initial student model to be trained;

[0014] Inputting at least one sample image into the student model to be trained for the t-th time and the pre-trained teacher model to obtain a prediction score output by the student model to be trained for the t-th time, the first type of attention map output by the student model to be trained for the t-th time, and the second type of attention map output by the pre-trained teacher model; t is an integer greater than or equal to 1;

[0015] Determining a loss value for the t-th time according to the prediction score output by the student model to be trained for the t-th time, the first type of attention map output by the student model to be trained for the t-th time, the second type of attention map output by the pre-trained teacher model, and the real label of the at least one sample image;

[0016] perform back propagation according to the loss value of the tth time and update parameters of the tth student model to be trained to obtain a t+1th student model to be trained, and iterate until a preset condition is met to stop training and obtain a trained student model to be converted;

[0017] replace a one-dimensional convolution kernel connected before the global pooling layer in the trained student model to be converted with a fully connected layer connected after the global pooling layer to obtain the trained student model.

[0018] In some embodiments, the determining of the loss value of the tth time according to the prediction score output by the tth student model to be trained, the first type of attention map output by the tth student model to be trained, the second type of attention map output by the pre-trained teacher model, and the real label of the at least one sample image comprises:

[0019] the determining of the knowledge distillation loss value of the tth time according to the first type of attention map output by the tth student model to be trained and the second type of attention map output by the pre-trained teacher model;

[0020] the determining of the cross-entropy loss value of the tth time according to the prediction score output by the tth student model to be trained and the real label of the at least one sample image;

[0021] the determining of the loss value of the tth time according to the knowledge distillation loss value of the tth time and the cross-entropy loss value of the tth time.

[0022] In some embodiments, the determining of the knowledge distillation loss value of the tth time according to the first type of attention map output by the tth student model to be trained and the second type of attention map output by the pre-trained teacher model comprises:

[0023] performing a pooling operation on the first type of attention map output by the tth student model to be trained to obtain first pooling features;

[0024] performing a pooling operation on the second type of attention map output by the pre-trained teacher model to obtain second pooling features;

[0025] performing L2 regularization processing on the first pooling features and the second pooling features respectively to obtain first processing features and second processing features;

[0026] taking a mean square error between the first processing features and the second processing features as the knowledge distillation loss value of the tth time.

[0027] In some embodiments, the determining the tth loss value according to the tth knowledge distillation loss value and the tth cross-entropy loss value comprises:

[0028] summing the tth knowledge distillation loss value and the tth cross-entropy loss value to obtain a summation value;

[0029] taking the summation value as the tth loss value.

[0030] In some embodiments, the obtaining the tth to-be-trained student model comprises:

[0031] when t is 1, performing initialization of the student model to obtain the initial to-be-trained student model.

[0032] The present application has the following beneficial technical effects:

[0033] Through the above method, the student model obtained by training can capture more recognition areas related to the target class in the graph, and improve the classification accuracy of the model while having high interpretability.

[0034] The present application will be further described in detail below in combination with the drawings and embodiments. BRIEF DESCRIPTION OF DRAWINGS

[0035] Figure 1 A flowchart of a model compression method based on class attention transmission provided by an embodiment of the present application;

[0036] Figure 2 An exemplary structure of the model proposed by the present application (Converted structure) and a conversion process diagram of the model structure from a normal model (Normoal CNN) provided by an embodiment of the present application;

[0037] Figure 3 An exemplary principle diagram for determining a knowledge distillation loss value provided by an embodiment of the present application;

[0038] Figure 4 A comparison diagram of the class attention map of the trained student model and the class attention map of the normal training model of an embodiment of the present application;

[0039] Figure 5 An exemplary training consumption and effect comparison diagram of a plurality of other methods and the method proposed by the present application provided by an embodiment of the present application. DETAILED DESCRIPTION

[0040] The present application will be further described in detail below in combination with specific embodiments, but the embodiments of the present application are not limited thereto.

[0041] In the description of the present application, the terms "first", "second" are only for descriptive purpose, and cannot be understood as indicating or implying relative importance or implicitly indicating the number of the indicated technical features. Therefore, the features defined with "first", "second" can explicitly or implicitly include one or more of the features. In the description of the present application, the meaning of "a plurality of" is two or more, unless otherwise explicitly specified.

[0042] In the description of the present application, the description of the terms "one embodiment", "some embodiments", "an example", "a specific example", or "some examples" and the like means that the specific features, structures, materials or characteristics described in connection with the embodiment or example are included in at least one embodiment or example of the present application. In the description of the present application, the illustrative description of the above terms is not necessarily directed to the same embodiment or example. Moreover, the specific features, structures, materials or characteristics described can be combined in any one or more embodiments or examples in a suitable manner. In addition, those skilled in the art can combine and combine different embodiments or examples described in the specification.

[0043] Although the present application is described herein in connection with various embodiments, those skilled in the art will understand and appreciate the disclosure of the disclosed embodiments upon reading the description of the drawings, the disclosure, and the appended claims. In the claims, the word "comprising" does not exclude other components or steps, and "a" or "one" does not exclude a plurality. A single processor or other unit can implement several functions listed in the claims. Some measures are described in mutually different dependent claims, but this does not mean that these measures cannot be combined to produce good results.

[0044] Figure 1 is a flowchart of the model compression method based on attention-like transmission provided by an embodiment of the present application, as shown in Figure 1 The method comprises the following steps:

[0045] S101, obtaining an image to be recognized.

[0046] S102, using a trained student model to recognize a target object in the image to be recognized to obtain the category of the target object; the trained student model is obtained by training a student model to be trained according to the prediction score output by the student model to be trained, the real label corresponding to the training sample, the first attention map output by the student model to be trained, and the second attention map output by the pre-trained teacher model; the fully connected layer connected after the global pooling layer in the student model to be trained and the teacher model is replaced by a one-dimensional convolution kernel connected before the global pooling layer.

[0047] Here, the student model and the teacher model are both models for performing target classification.

[0048] Illustratively, the student model to be trained and the pre-trained teacher model both include a convolutional layer, a one-dimensional convolutional kernel, and a global pooling layer; wherein the convolutional layer is configured to extract features of an input image to obtain a feature map of the input image; the one-dimensional convolutional kernel is configured to generate a class attention map according to the feature map; and the global pooling layer is configured to convert the class attention map into a prediction score.

[0049] Illustratively, the trained student model includes a convolutional layer, a global pooling layer, and a fully connected layer.

[0050] As shown in Figure 2 In an image classification task, a mainstream model usually uses a convolutional neural network (CNN) to extract features first, and then a global pooling (GAP) is performed on a plurality of feature maps generated by the CNN to obtain prediction scores, which are then input into a simple fully connected (FC) layer to perform a classification task. However, in the present application, when training a student model, a fully connected layer is converted into a one-dimensional convolutional kernel (1x1 Conv), so that a class activation map (CAM) (for example, K class activation maps in Figure 2 ) can be obtained during the forward propagation of the CNN.

[0051] In some embodiments, before S102, the following steps are further included:

[0052] S201, obtaining a plurality of sample images with real labels; the real labels representing real classes of objects in the sample images.

[0053] S202, obtaining a t-th student model to be trained and a pre-trained teacher model; wherein when t is 1, the student model to be trained is an initial student model to be trained.

[0054] Here, when t is 1, the student model is initialized to obtain the initial student model to be trained.

[0055] S203, inputting at least one sample image into the t-th student model to be trained and the pre-trained teacher model to obtain a prediction score output by the t-th student model to be trained, a first class attention map output by the t-th student model to be trained, and a second class attention map output by the pre-trained teacher model; t is an integer greater than or equal to 1.

[0056] Specifically, at each training time, one or more sample images can be randomly selected, and the selected sample images are input into the student model to be trained this time to obtain the class attention map (referred to as the first class attention map) output by the student model to be trained this time and the prediction score output by the student model to be trained this time; and the selected sample images are input into the pre-trained teacher model to obtain the class attention map (referred to as the second class attention map) output by the pre-trained teacher model.

[0057] S204, according to the prediction score output by the student model to be trained t times, the first class attention map output by the student model to be trained t times, the second class attention map output by the pre-trained teacher model, and the true label of the at least one sample image, determine the loss value of the tth time.

[0058] Here, the knowledge distillation loss value of the tth time can be determined according to the first class attention map output by the student model to be trained t times and the second class attention map output by the pre-trained teacher model; the cross-entropy loss value of the tth time can be determined according to the prediction score output by the student model to be trained t times and the true label of the at least one sample image; the sum of the knowledge distillation loss value of the tth time and the cross-entropy loss value of the tth time is summed to obtain a sum value; and the sum value is taken as the loss value of the tth time.

[0059] Specifically, the principle of determining the knowledge distillation loss value of the tth time is as follows: the first class attention map output by the student model to be trained t times is subjected to a pooling operation to obtain a first pooling feature; the second class attention map output by the pre-trained teacher model is subjected to a pooling operation to obtain a second pooling feature; the first pooling feature and the second pooling feature are subjected to L2 regularization processing respectively to obtain a first processed feature and a second processed feature; and the mean square error between the first processed feature and the second processed feature is taken as the knowledge distillation loss value of the tth time.

[0060] Exemplarily, Figure 3 The principle diagram for determining the knowledge distillation loss value is shown in FIG. 1. Figure 3As shown, the same sample image is respectively input into the pre-trained teacher model and the student model to be trained, and the class attention map Transferred CAMs output by the pre-trained teacher model and the class attention map CAMs output by the student model to be trained are obtained. After the Transferred CAMs and the CAMs are respectively subjected to pooling processing, the pooling features corresponding to the pre-trained teacher model with the size of W*H*K (hereinafter referred to as first pooling features) and the pooling features corresponding to the student model to be trained with the size of W*H*K (hereinafter referred to as second pooling features) are obtained. After the first pooling features and the second pooling features are respectively subjected to normalization, the normalized features corresponding to the pre-trained teacher model (hereinafter referred to as first processing features) and the normalized features corresponding to the student model to be trained (hereinafter referred to as second processing features) are obtained. The mean square error between the first processing features and the second processing features is calculated, and the error is taken as a knowledge distillation loss value (CATloss).

[0061] S205, according to the loss value of the tth time, the parameters of the tth student model to be trained are updated by back propagation, and the t+1th student model to be trained is obtained. Thus, iteration is performed until the preset condition is met to stop training, and the trained student model to be converted is obtained.

[0062] Here, when the training times reach the preset times, or the loss value of a certain time reaches the preset threshold, it is determined that the preset condition is met. At this time, the training can be stopped, and the student model obtained by the last training is taken as the trained student model to be converted.

[0063] S206, the one-dimensional convolution kernel connected before the global pooling layer in the trained student model to be converted is replaced by the full connection layer connected after the global pooling layer, and the trained student model is obtained.

[0064] Here, since the structure of the trained student model to be converted obtained is as shown in Convertedstructure in Figure 2 Therefore, by replacing the one-dimensional convolution kernel connected before the global pooling layer in the trained student model to be converted with the full connection layer connected after the global pooling layer, the trained student model can be obtained, and the conversion does not change the prediction score output by the model.

[0065] Through the above method, the student model trained by the present application can capture more recognition areas related to the target class in the image, and improve the classification accuracy of the model while having high interpretability.

[0066] The technical effects achieved by the embodiments of the present application will be further described through experimental data.

[0067] Figure 4 Comparison of class attention maps of the trained student model of the present method (CAT-KD) and the class attention maps of the ordinary trained model. Figure 4 The three figures in the first row of (a) are visualizations of the class attention maps of the ordinary trained model, Figure 4 The three figures in the second row of (b) are visualizations of the class attention maps of the trained student model of the present method. As shown in (b), Figure 4 the trained student model of the present application can capture more recognition areas related to the target class in the figure.

[0068] Table 1 below is a comparison of the effects of the present method (CAT-KD) and other knowledge distillation methods (KD, DKD, CRD, OFD, FitNet, RKD, Review KD, AT) on the CIFAR100 dataset. Obviously, the present method has better effect. Table 2 is a comparison of the results obtained by transfer learning on the models trained by multiple other methods and the model trained by the present method. Obviously, the model trained by the present method has better generalization.

[0069]

[0070]

[0071] Table 1

[0072]

[0073] Table 2

[0074] Figure 5 The training consumption and effect comparison diagram of multiple other methods and the present method. Obviously, the present method is the best.

[0075] The above is a further detailed description of the present application in combination with specific preferred embodiments, and the specific implementation of the present application cannot be limited to these descriptions. For ordinary skilled persons in the technical field to which the present application belongs, without departing from the concept of the present application, a number of simple deductions or substitutions can be made, which should be considered as falling within the protection scope of the present application.

Claims

1. A model compression method based on attention-based transfer, characterized in that, include: Acquire the image to be recognized; The trained student model is used to identify the target object in the image to be identified, and the category of the target object is obtained; The trained student model is obtained by training the student model to be trained using the predicted score output by the training samples, the true labels corresponding to the training samples, the first type of attention map output by the student model to be trained itself, and the second type of attention map output by the pre-trained teacher model. Both the student model to be trained and the teacher model include: a convolutional layer, a one-dimensional convolutional kernel, and a global pooling layer. The fully connected layer after the global pooling layer in the student model to be trained and the teacher model is replaced with a one-dimensional convolutional kernel connected before the global pooling layer. Before using a trained student model to identify target objects in the image to be identified and obtaining the category of the target objects, the method further includes: Acquire multiple sample images with real labels; the real labels represent the true categories of objects in the sample images; Obtain the student model to be trained for the tth time, and the pre-trained teacher model; wherein, when t is 1, the student model to be trained is the initial student model to be trained; At least one sample image is input into the student model to be trained for the tth time and the pre-trained teacher model to obtain the predicted score output by the student model to be trained for the tth time, the first type of attention map output by the student model to be trained for the tth time, and the second type of attention map output by the pre-trained teacher model; t is an integer greater than or equal to 1; The loss value for the tth iteration is determined based on the predicted score output by the student model to be trained for the tth iteration, the first type of attention map output by the student model to be trained for the tth iteration, the second type of attention map output by the pre-trained teacher model, and the true label of the at least one sample image. Backpropagation is performed based on the loss value of the tth iteration to update the parameters of the student model to be trained in the tth iteration, and the student model to be trained in the (t+1)th iteration is obtained. This process is repeated until the preset conditions are met and training stops, resulting in a trained student model to be converted. The one-dimensional convolutional kernel connected before the global pooling layer in the trained student model to be converted is replaced with a fully connected layer connected after the global pooling layer to obtain the trained student model.

2. The model compression method based on attention transfer as described in claim 1, characterized in that, Both the student model and the teacher model are models used for target classification.

3. The model compression method based on attention transfer as described in claim 1, characterized in that, The convolutional layer is used to extract features from the input image to obtain a feature map of the input image; the one-dimensional convolutional kernel is used to generate a class attention map based on the feature map; and the global pooling layer is used to convert the class attention map into a prediction score.

4. The model compression method based on attention transfer according to claim 1, characterized in that, The step of determining the loss value for the t-th iteration based on the predicted score output by the student model to be trained for the t-th iteration, the first type of attention map output by the student model to be trained for the t-th iteration, the second type of attention map output by the pre-trained teacher model, and the true labels of the at least one sample image includes: The knowledge distillation loss value for the tth time is determined based on the first type of attention map output by the student model to be trained for the tth time and the second type of attention map output by the pre-trained teacher model. The cross-entropy loss value for the tth iteration is determined based on the predicted score output by the student model to be trained for the tth iteration and the true label of the at least one sample image. The loss value for the tth time is determined based on the knowledge distillation loss value for the tth time and the cross-entropy loss value for the tth time.

5. The model compression method based on attention transfer according to claim 4, characterized in that, Determining the knowledge distillation loss value for the t-th time based on the first type of attention map output by the student model to be trained for the t-th time and the second type of attention map output by the pre-trained teacher model includes: The first type of attention map output by the student model to be trained for the tth time is subjected to a pooling operation to obtain the first pooling feature; The second type of attention map output by the pre-trained teacher model is subjected to a pooling operation to obtain the second pooling feature; The first pooling feature and the second pooling feature are subjected to L2 regularization to obtain the first processed feature and the second processed feature. The mean square difference between the first processed feature and the second processed feature is used as the knowledge distillation loss value for the t-th time.

6. The model compression method based on attention transfer according to claim 4, characterized in that, Determining the loss value for the t-th iteration based on the knowledge distillation loss value and the cross-entropy loss value for the t-th iteration includes: The knowledge distillation loss value of the t-th iteration and the cross-entropy loss value of the t-th iteration are summed to obtain the sum value; The summation value is used as the loss value for the t-th iteration.

7. The model compression method based on attention transfer according to claim 4, characterized in that, The process of obtaining the student model to be trained for the t-th time includes: When t is 1, the student model is initialized to obtain the initial student model to be trained.

Images