An image classification method based on knowledge distillation

Abstract

The application discloses a kind of based on knowledge distillation's image classification method, comprising: (1) constructing a deep neural network as first depth model;(2) the first depth model constructed in image training set is trained;(3) extracting the network backbone part in the first depth model trained, join auxiliary classification head after the output of each layer, randomly initialize auxiliary classification head, constitute second depth model;(4) freeze network backbone part, train auxiliary classification head;(5) construct a depth neural network including each layer auxiliary classification head, carry out random initialization, constitute third depth model;(6) train third depth model layer by layer, utilize the output of each layer auxiliary classification head of second depth model as supervision;(7) extract the backbone part of third depth model and last layer auxiliary classification head, carry out image classification.Using the application, the optimization ability of training model can be improved, and the effect of image classification is improved.

Description

TECHNICAL FIELD

[0001] The application belongs to the field of computer vision and image processing, and particularly relates to an image classification method based on knowledge distillation. BACKGROUND

[0002] The development of deep learning has promoted the rapid breakthrough and significant improvement of multiple branch tasks in the field of computer vision including computer vision and natural language processing. The standard training mode of deep neural network is end-to-end training, that is, all parameters in the deep neural network are optimized in the same training step, which requires loading the entire model into the video memory for gradient optimization. This end-to-end training mode makes the memory cost of training extremely high, and this shortcoming is particularly prominent today when large model applications are gradually widespread.

[0003] Another training mode relative to end-to-end training is layer-by-layer training, which was first applied to the pre-training of neural network by Yoshua Bengio et al. in 2006 in the work Advances in neural information processing systems Greedy Layer-wise Training of Deep Networks. The work divides the optimization problem of the network into several sub-problems, and solves each sub-problem one by one, only needing to train one layer of the network in each sub-problem, that is, only part of the parameters of the network need to be loaded into the video memory, so the demand for video memory in the training process can be significantly reduced. Eugene Belilovsky et al. introduced the training scheme of layer-by-layer training on large-scale data sets such as ImageNet in the work Greedy Layerwise Learning Can Scale to ImageNet published in the 2019 International Conference on Machine Learning.

[0004] However, the effectiveness of layer-by-layer training is not fully understood, and its effect is also flawed, which cannot be compared with the effect of end-to-end training. By carefully analyzing and comparing the processes of layer-by-layer training and end-to-end training, it can be found that in the training of each layer of layer-by-layer training, the training target of each layer is the same as that of end-to-end training. For end-to-end training, the training target is to minimize the cross-entropy loss of the network output probability and the real label. Layer-by-layer training also adopts this training mode, and in the training of each layer, the cross-entropy loss of the output probability of the layer and the real label is minimized. However, the difference between layer-by-layer training and end-to-end is that we expect different layers in layer-by-layer training to learn different feature representations, so we need to have differentiated supervision in the spatial dimension. However, this training mode with the same supervision and training target for each layer causes the network to lose differentiated supervision in the spatial dimension, and there is no mechanism in the training to explicitly guide the neural network to focus on the information in a certain region in space. SUMMARY

[0005] The present application provides a knowledge distillation-based image classification method, which can promote differentiated supervision in the spatial dimension during training, improve the optimization ability of the trained model, reduce the memory requirement in training, and improve the effect of image classification.

[0006] A knowledge distillation-based image classification method, comprising the following steps:

[0007] (1) constructing a deep neural network as a first deep model;

[0008] (2) training the constructed first deep model on an image training set;

[0009] (3) extracting the network backbone part of the trained first deep model, adding an auxiliary classification head after the output of each layer, randomly initializing the auxiliary classification head, and constructing a second deep model;

[0010] (4) for the second deep model, freezing the network backbone part, training the auxiliary classification head, and training the same number of times as the first deep model;

[0011] (5) constructing a deep neural network containing an auxiliary classification head in each layer, and randomly initializing to form a third deep model;

[0012] (6) training the third deep model layer by layer, using the output of the auxiliary classification head of each layer of the trained second deep model as a soft label to supervise the third deep model, and training the same number of times as the first deep model;

[0013] (7) extracting the backbone part of the third deep model and the auxiliary classification head of the last layer to form a fourth deep model, and using the fourth deep model to analyze the to-be-tested picture and predict the image category.

[0014] The present application firstly constructs a deep neural network, trains it on a large-scale dataset, then trains an auxiliary classification head for each layer of the deep neural network, and then trains a deep neural network containing an auxiliary classification head for each layer layer by layer using the output of the auxiliary classification head of the neural network as supervision. The backbone of the deep neural network trained layer by layer and the auxiliary classification head of the last layer are used as the final model, and the model is used to predict the category of the input image. This method can promote the deep neural network to learn the potential task of "feature expression in each layer space", thereby producing a layer-by-layer training model with better optimization effect.

[0015] In step (1), the structure of the first deep model is as follows:

[0016] The deep neural network has J layers, and the input representation is X, and the initial feature representation is The features of the shallower layer are taken as input to further extract features, and the new features are passed to the deeper layer to obtain the final feature representation X J ;

[0017] For the final feature representation X J , it is first processed by the global average pooling layer g, and the spatial dimension of X J is reduced, thereby generating a single vector g(X J ); The training head h ω parameterized by ω converts the shape of g(X J ) and calculates the loss value, as follows:

[0018]

[0019] In the formula, is the loss function, and y is the label.

[0020] In step (2), the image training set uses ImageNet.

[0021] In step (2), end-to-end training is used, and all parameters of the first deep model are optimized:

[0022]

[0023] In the formula, J represents the number of network layers of the neural network, and the parameters of the jth network module are represented by θ j , and the trained parameters are represented by .

[0024] In step (3), the size of the jth layer feature X j is [H j , W j], in order to differentiate supervision in the spatial dimension, the image label y is copied multiple times to extend to [H j , W j ] size, denoted by Y; the jth layer feature X j needs to be converted by the auxiliary classification head h to match the size of Y.

[0025] The auxiliary classification head of layer j is denoted as

[0026] In step (4), the optimization objective of the auxiliary classification head of each layer is denoted as:

[0027] For the first layer auxiliary classification head:

[0028] For the second layer auxiliary classification head:

[0029] …

[0030] For the Jth layer auxiliary classification head:

[0031] The auxiliary classification head of each layer is trained in parallel, and the trained parameters are denoted by .

[0032] The specific process of step (6) is:

[0033] (6-1) input a batch of data into the second deep model;

[0034] (6-2) the output of the auxiliary classification head of the jth layer of the second deep model is used as the soft label:

[0035]

[0036] (6-3) train the jth layer of the third deep model, and only the network backbone and auxiliary classification head of the jth layer in the third deep model are updated:

[0037]

[0038] (6-4) sequentially perform (6-2) and (6-3) on each layer to complete the training.

[0039] Compared with the prior art, the present application has the following beneficial effects:

[0040] 1. The present application is based on the understanding of the optimization of deep neural networks, combined with the knowledge distillation technology, and designs a new type of layer-by-layer training method, which highlights the differentiated supervision of deep neural networks in the spatial dimension, making it easier to learn effective feature expression in the spatial dimension when training a single layer, thereby improving the image classification effect.

[0041] 2. This invention has strong applicability and can be directly and effectively deployed in most existing deep neural network models by replacing the network backbone, without affecting other methods and processes. Attached Figure Description

[0042] Figure 1 This is a schematic diagram of the image classification method based on knowledge distillation according to the present invention;

[0043] Figure 2 This is a schematic diagram of a layer-by-layer training mode based on knowledge distillation in an embodiment of the present invention. Detailed Implementation

[0044] The present invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be noted that the embodiments described below are intended to facilitate the understanding of the present invention and do not constitute any limitation thereof.

[0045] like Figure 1 As shown, an image classification method based on knowledge distillation includes the following steps:

[0046] S01, construct a deep neural network as the first deep model (teacher model).

[0047] This deep neural network has J layers. Let the input be X, and the initial features be denoted as... After the feedforward propagation through layer j, the initial feature X0 is transformed into X j Each network layer can be any predefined network module, constructed from basic operations such as linear transformations, white attention, batch normalization, and activation functions. The j-th layer... The parameters are determined by θ j This means that features from shallower layers are used as input to further extract features, and the new features are then passed to deeper layers:

[0048]

[0049] In visual tasks, X j It is a feature map whose shape consists of three dimensions, including two spatial dimensions (feature map height and width) and one channel dimension.

[0050] Ultimately, X represents J It will first be processed by the global average pooling layer g, during which X... J The spatial dimension will be reduced, resulting in a single vector g(X). J For self-attention models, the global average pooling operation in the final representation is replaced by a process for extracting class-labeled features. Finally, the training head h is parameterized by ω. ω g(X) J Convert to a suitable shape and calculate the loss value:

[0051]

[0052] is the loss function, y is the label. In most cases, h ω only contains a linear transformation.

[0053] S02, the constructed deep neural network is trained on a large-scale dataset. Here, ImageNet is used for training, which is the first super large-scale image recognition dataset published by Jia Deng et al. in the paper “Imagenet: A large-scale hierarchical image database” at the Conference on Computer Vision and Pattern Recognition 2009. The dataset is derived from the Internet and contains pictures of 1000 categories of objects, with an image size of about 256*256, and more than 1000 training pictures for each category. The training set contains 1281167 pictures, and the verification set contains 50000 pictures.

[0054] End-to-end training is used, and all parameters in the model are optimized at the same time:

[0055]

[0056] The trained parameters are represented by .

[0057] S03, the network backbone part in step S02 is extracted, and an auxiliary classification head is added after the output of each layer in the model. The parameters of the auxiliary classification head are randomly initialized.

[0058] The size of the jth layer feature X j is [H j , W j ], in order to perform differential supervision in the spatial dimension, the image label y is copied and expanded to the size of [H j , W j ], denoted by Y. The jth layer feature X j needs to be converted through the auxiliary classification head h to match the size of Y. The auxiliary classification head of layer j is denoted as

[0059] S04, the parameters of the network backbone part of the teacher model are frozen, and the auxiliary classification head is trained. The training hyperparameters and the number of times are controlled to be the same as the network backbone training.

[0060] The optimization objective of the auxiliary classification head of each layer is represented as:

[0061] For the first layer auxiliary classification head:

[0062] For the second layer auxiliary classification head:

[0063] …

[0064] For the Jth layer auxiliary classification head:

[0065] The auxiliary classification heads of each layer can be trained in parallel and do not have a dependency relationship, and do not need to be trained sequentially.

[0066] The trained parameters are denoted by .

[0067] S05, a deep neural network containing auxiliary classification heads of each layer is constructed, and is randomly initialized as a student model.

[0068] S06, the student model is trained layer by layer, and the output of the teacher model is used as a soft label to supervise the student model, and the training hyperparameters and times are controlled to be the same as the teacher model training.

[0069] The specific process is as follows:

[0070] (6-1) input a batch of data into the teacher network

[0071] (6-2) the output of the auxiliary classification head of the jth layer of the teacher network is used as a soft label:

[0072]

[0073] (6-3) train the jth layer student network, only the network backbone and the auxiliary classification head of the jth layer in the student network are updated:

[0074]

[0075] (6-4) sequentially perform (6-2) and (6-3) on each layer to complete the training.

[0076] As shown in Figure 2 , in the process of training the student model layer by layer, all parameters of the teacher model are frozen. When training the jth layer student network, only the network backbone and the corresponding auxiliary classification head of the jth layer are trainable, and the parameters of the remaining layers are frozen. The output of the auxiliary classification head of the jth layer of the teacher network is used as a soft label to perform knowledge distillation and train the jth layer student network. The J layers of the student model are sequentially trained from step 1 to step J.

[0077] S07, extract the backbone part of the student model obtained by the layer-by-layer training and the auxiliary classification head of the last layer as the final model, and use the final model to perform an image classification task.

[0078] To prove the effectiveness of the method of the present application, tests were carried out on the ImageNet dataset. The models were ViT-T / 16, ViT-S / 16, ViT-B / 16, ResNet-18, ResNet-34, and ResNet-50. For the image classification task, the Top-1 accuracy and Top-5 accuracy in the ImageNet were used as evaluation indexes, and the results are shown in Table 1.

[0079] Table 1

[0080]

[0081] Comparing the effect of the layer-by-layer training method based on knowledge distillation (as shown in SLD in Table 1) and the effect of the end-to-end training method (as shown in teacher in Table 1), the layer-by-layer training method based on knowledge distillation can improve the image classification effect to the level of the end-to-end training on visual self-attention models and ResNet models of different depths.

[0082] To prove the effectiveness of the method of the present application, tests were carried out on the ImageNet dataset. The models were ViT-T / 16, ViT-S / 16, ViT-B / 16, ResNet-18, ResNet-34, and ResNet-50. For the image classification task, the Top-1 accuracy and Top-5 accuracy in the ImageNet were used as evaluation indexes, and the results are shown in Table 1.

[0083] The above embodiments describe the technical solutions and beneficial effects of the present application in detail. It should be understood that the above description is only a specific embodiment of the present application and is not intended to limit the present application. Any modification, supplement and equivalent replacement made within the principle range of the present application shall be included in the protection scope of the present application.

Claims

1. A method for image classification based on knowledge distillation, characterized in that, The method comprises the following steps: (1) constructing a deep neural network as a first deep model; (2) training the constructed first deep model on an image training set; (3) extracting a network backbone part in the trained first deep model, adding an auxiliary classification head after the output of each layer, and randomly initializing the auxiliary classification head to form a second deep model; (4) for the second deep model, freezing the network backbone part, training the auxiliary classification head, and training the same number of times as the first deep model; (5) constructing a deep neural network containing an auxiliary classification head of each layer, performing random initialization, and forming a third deep model; (6) training the third deep model layer by layer, using the output of each layer of the trained second deep model as a soft label to supervise the third deep model, and training the same number of times as the first deep model; The specific process is as follows: (6-1) input a batch of data into the second deep model; (6-2) the output of the auxiliary classification head of the jth layer of the second deep model is used as a soft label: (6-3) training the jth layer of the third deep model, only the network backbone and the auxiliary classification head of the jth layer in the third deep model are updated: (6-4) sequentially perform (6-2) and (6-3) on each layer to complete the training; (7) extracting the backbone part of the third deep model and the auxiliary classification head of the last layer to form a fourth deep model, and using the fourth deep model to analyze the to-be-tested picture and predict the image category.

2. The knowledge distillation based image classification method of claim 1, wherein, In step (1), the structure of the first deep model is as follows: A deep neural network has J layers, takes an input representation X, and takes an initial feature representation Further extracts features taking the features of the shallower layers as input, and passes the new features to the deeper layers to get a final feature representation X J ; For the final feature representation X J , it is first processed by a global average pooling layer g, X J , the spatial dimensions are reduced, resulting in a single vector g(X J ); a trained head h ω parameterized by ω converts g(X J ) to shape and computes a loss value, as follows: In the formula, is a loss function, y is a label. 3.The knowledge distillation based image classification method of claim 1, wherein, In step (2), the image training set adopts ImageNet.

4. The knowledge distillation based image classification method of claim 2, wherein, In step (2), end-to-end training is adopted, and all parameters of the first deep model are optimized: In the formula, J represents the number of network layers of the neural network, the parameter of the jth layer network module is represented by θ , and the trained parameter is represented by θ j . ​ 5. The knowledge distillation based image classification method according to claim 4, characterized in that, In step (3), the j-th layer feature X j has size [H j ,W j ] and, in order to supervise the differentiation in the spatial dimension, the image label y is replicated multiple times to expand to size [H j ,W j ] denoted by Y; the j-th layer feature X j needs to be transformed by an auxiliary classification head h to match the size of Y; the auxiliary classification head for layer j is denoted as 6. The knowledge distillation based image classification method according to claim 5, characterized in that, In step (4), the optimization target of each layer of the auxiliary classification head is represented as: For the first layer auxiliary classification head: For the second layer auxiliary classification head: … To the J-th auxiliary classification head: The auxiliary classification heads of each layer are trained in parallel, and the trained parameters are used to represent.

Images