An image classification self-distillation method based on mask feature reconstruction

Abstract

The application discloses an image classification self-distillation method based on mask feature reconstruction, which comprises the following steps: dividing a preprocessed input image into a training set and a verification set, and enhancing the training set; inputting the preprocessed image into a residual neural network to obtain features output by shallow and deep layers of a residual block and a prediction output of the residual neural network, and performing down-sampling on the features by using convolution; performing element dot product on the down-sampled features and a random mask to obtain mask features; reconstructing the mask features by using a convolution-based generator and a perception-based generator, and obtaining a prediction output of the reconstructed features; combining a classification loss and a reconstruction loss to construct a self-distillation loss function; inputting images in the enhanced training set into the residual neural network, performing back propagation training according to the self-distillation loss function, verifying by using the verification set, and saving the best network weight and the network weight of each training round. The application realizes the improvement of network performance and maintains high training efficiency.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of image classification and self-knowledge distillation, in particular to an image classification self-distillation method based on mask feature reconstruction. BACKGROUND

[0002] In recent years, deep neural networks have achieved great success in image classification, object detection and image segmentation due to their ability to operate on large parameters, but it is difficult to deploy high-capacity models on resource-constrained devices. Model compression techniques aim to maintain high accuracy while reducing the parameter amount and computational amount of high-capacity models. Among them, knowledge distillation is a simple and effective model compression method, which forces a small model to imitate the knowledge of a large model so that the small model can achieve the same accuracy as the large model. However, the traditional knowledge distillation algorithm is divided into two stages, first a teacher model needs to be pre-trained, and then the teacher model is used to distill the student model. Although this can effectively improve the performance of the student model, a good and careful teacher design is challenging, and pre-training the teacher requires more training time. To solve this problem, self-learning-based distillation methods have emerged.

[0003] Although the self-distillation algorithm shows potential in improving the performance of the student model, the existing self-distillation algorithm still has the following main shortcomings: (1) Due to the lack of teacher guidance, the existing self-distillation algorithm often uses complex multi-branch complex structures to obtain additional supervision signals, which increases the complexity of the training process to some extent and also increases the training time. (2) Most existing auxiliary structure-based methods use block stacking to increase the trainability of shallow layers, ignoring further optimization of shallow features. SUMMARY

[0004] The purpose of the present application is to overcome the shortcomings and shortcomings of the prior art, and to provide an image classification self-distillation method based on mask feature reconstruction, which realizes the self-distillation strategy through the mask shallow feature reconstruction method, simplifies the self-distillation process, increases the trainability of the shallow layer while optimizing the features, improves the network performance, and maintains high training efficiency.

[0005] To achieve the above purpose, the technical scheme provided by the present application is: an image classification self-distillation method based on mask feature reconstruction, comprising the following steps:

[0006] Step S1: After data preprocessing based on image mixing, the input image is divided into training set and validation set according to the proportion, and the training set is subjected to data enhancement in two ways of cropping and rotating;

[0007] Step S2: Input the preprocessed image into the residual neural network, which contains 4 residual blocks, which are divided into residual block 1, residual block 2, residual block 3 and residual block 4 in sequence. Residual block 1, residual block 2 and residual block 3 are considered to be shallow stages due to the small number of residual structure layers, and residual block 4 is considered to be a deep stage due to the large number of residual structure layers. The features output by residual block 2 and residual block 4 and the predicted output of the residual neural network are obtained through forward propagation, and the features output by residual block 2 are downsampled using convolution with a step size of 4;

[0008] Step S3: Design a random mask and perform element-wise product of the downsampled features and the random mask to obtain the mask features;

[0009] Step S4: Use a convolution-based generator and a perceptron-based generator to reconstruct the mask features, and use an average pooling operation, a linear classifier, and a multi-classification function softmax to obtain the predicted output of the reconstructed features;

[0010] Step S5: Use the cross entropy function to calculate the difference between the predicted output of the residual neural network and the true label, which is called classification loss, and use the KL divergence to measure the difference between the predicted output of the residual neural network in the current training round and the predicted output of the residual neural network history model in the previous training round, which is called reconstruction loss. By combining the classification loss and the reconstruction loss, a self-distillation loss function is constructed;

[0011] Step S6: Input the enhanced images in the training set into the residual neural network, perform backpropagation training according to the constructed self-distillation loss function, and use the validation set to verify the training effect of the network until the residual neural network converges, and save the optimal network weights and the network weights of each training round.

[0012] Furthermore, in step S1, two training images x1 and x2 are randomly selected and mixed using linear interpolation:

[0013] x m =λx1+(1-λ)x2,y m =λy1+(1-λ)y2

[0014] Where x m Represents the mixed image after linear interpolation, y1, y2 and y m They represent the true labels corresponding to the two images and the true labels corresponding to the mixed image, and λ is the image mixing ratio.

[0015] Furthermore, in step S2, the features output by residual block 2 are downsampled by convolution with a step size of 4 and a convolution kernel of 4 to achieve alignment with the features output by residual block 4, so as to reduce information loss in the feature map.

[0016] Further, the step S3 comprises the following steps:

[0017] S31: design of random mask:

[0018]

[0019] In the formula, M h,v represents a random mask, R h,v represents a random number between 0 and 1, h and v represent the horizontal and vertical coordinates corresponding to the feature map respectively, and γ represents the mask rate.

[0020] S32: element dot product of the down-sampled feature and the designed random mask, so as to obtain the masked feature:

[0021]

[0022] In the formula, F(x m ) and F'(x m ) represent the features before and after masking respectively, represents element multiplication.

[0023] Further, the step S4 comprises the following steps:

[0024] S41: design of the generator based on convolution, containing two 3x3 convolution layers, a GELU activation function layer and a batch normalization layer;

[0025] S42: design of the generator based on perception, containing two 1x1 convolution layers, a GELU activation function layer and a layer normalization layer;

[0026] S43: sequentially use the generator based on convolution and the generator based on perception for feature generation, and the reconstruction mode is as follows:

[0027] F"(x m ) = θ MLP (θ conv (F'(x m ))

[0028] In the formula, F"(x m ) represents the reconstructed feature, θ conv and θ MLP represent the generator based on convolution and the generator based on perception respectively.

[0029] S44: use the average pooling operation, the linear classifier and the multi-classification function softmax to obtain the probability prediction output corresponding to F"(x m ).

[0030] Further, the step S5 comprises the following steps:

[0031] S51: Calculate the classification loss L ce The difference between the predicted output of the residual neural network and the true label is calculated by using the cross-entropy function, which is calculated as follows:

[0032]

[0033] In the formula, K represents the number of categories in the training set, y m , p k represents the probability prediction of the kth category of the residual neural network;

[0034] S52: Calculate the reconstruction loss L rec1 The difference between the predicted output of the reconstruction feature and the predicted output of the residual neural network is measured by using the KL divergence, which is calculated as follows:

[0035]

[0036] In the formula, p k ′ represents the probability prediction of the kth category corresponding to the reconstruction feature after the average pooling operation, the linear classifier and the multi-classification function softmax;

[0037] S53: Calculate the reconstruction loss L rec2 After each round of training, the model is saved, and the model of the last round is loaded to predict the images x1 and x2 before mixing respectively to obtain the corresponding probability predictions and The difference between the predicted output of the reconstruction feature and is measured by using the KL divergence, which is calculated as follows:

[0038]

[0039] S54: The self-distillation loss function constructed contains the classification loss and two reconstruction losses:

[0040] L total = L ce + α(L rec1 + L rec2 )

[0041] In the formula, L total represents the self-distillation loss function, and α is a hyperparameter for balancing the classification loss and the reconstruction loss.

[0042] Further, in step S6, the batch size of the training is set to 128, the image is scaled to 32*32 size input into the residual neural network, the number of training iterations is set to 200, starting from 5 training epochs of linear warm-up of the initial learning rate from 0 to 0.1, multiplying the learning rate by 0.1 at the iteration times of 105 and 155 respectively, using the SGD optimizer; the validation set is data that does not contain training images, a total of 10000 images; when training, stop training when the maximum number of iterations is reached, and save the weight file corresponding to the best accuracy and the network weight of each training epoch.

[0043] Further, the residual neural network is a ResNet18 residual neural network.

[0044] Compared with the prior art, the present application has the following advantages and beneficial effects:

[0045] The present application only designs a single-branch reconstruction auxiliary structure, avoids the complexity of training caused by introducing a multi-branch auxiliary module, increases the trainability of the shallow layer, simplifies the self-distillation process, and improves the training efficiency of the network. In view of the problem that the existing algorithm often ignores the optimization of shallow features, the mask feature reconstruction method proposed in the present application reduces the background information in the shallow features, thereby optimizing the shallow features. BRIEF DESCRIPTION OF DRAWINGS

[0046] Figure 1 The network architecture diagram of one specific embodiment of the present application.

[0047] Figure 2 The architecture diagram of the convolution-based generator and the perception-based generator. DETAILED DESCRIPTION

[0048] The present application will be further described in detail below in conjunction with the embodiments and drawings, but the embodiments of the present application are not limited thereto.

[0049] In this embodiment, CIFAR100 is used as the image classification training data set, which is composed of a total of 60,000 pictures of 100 classification categories, of which 50,000 are training sets and 10,000 are validation sets.

[0050] As shown in Figure 1 and Figure 2 The present application discloses an image classification self-distillation method based on mask feature reconstruction, comprising the following steps:

[0051] Step S1: After data preprocessing based on image mixing, the input image is divided into training set and validation set according to the proportion, and the training set is subjected to data enhancement in two ways of cropping and rotating;

[0052] In one specific embodiment, CIFAR100 is used as the data set, and the training set and the validation set are divided in proportion, wherein the training set contains 50,000 images and the validation set contains 10,000 images.

[0053] Two training images x1 and x2 are randomly selected, and linear interpolation is used to mix the two images:

[0054] x m =λx1+(1-λ)x2,y m =λy1+(1-λ)y2

[0055] In the formula, x m represents the mixed image after linear interpolation, y1, y2 and y m respectively represent the true labels corresponding to the two images and the true label corresponding to the mixed image, and λ is the image mixing ratio, which is set to 0.4.

[0056] Step S2: input the preprocessed image into the residual neural network, obtain the features output by the shallow stage and the deep stage of the residual block and the prediction output of the residual neural network through forward propagation, and use convolution with a step size of 4 to downsample the shallow features, which specifically includes the following steps:

[0057] S21: referring to the ResNet18 residual neural network shown in Figure 1 , which contains 4 residual blocks, which are sequentially divided into residual block 1, residual block 2, residual block 3 and residual block 4, wherein residual block 1, residual block 2 and residual block 3 are regarded as shallow stages due to the small number of residual structure layers, and residual block 4 is regarded as a deep stage due to the large number of residual structure layers, the features output by residual block 2 and residual block 4 and the prediction output of the residual neural network are obtained through forward propagation, and the features output by residual block 2 are downsampled using convolution with a step size of 4;

[0058] S22: downsample the residual block 2 using convolution with a step size of 4 and a convolution kernel of 4 to align the output features of the residual block 4, and the convolution-based downsampling method can effectively reduce the information loss in the feature map, compared with the pooling-based downsampling method, which can reduce the information loss in the feature map, thereby improving the accuracy of reconstruction. At the same time, mask feature reconstruction at the shallow stage can enhance the trainability of the shallow stage.

[0059] Step S3: design a random mask, and perform element dot product on the downsampled features and the random mask to obtain the mask features, which specifically includes the following steps:

[0060] S31: random mask strategy:

[0061]

[0062] where M h,v represents a random mask, R h,v represents a random number between 0 and 1, h and v represent the horizontal and vertical coordinates corresponding to the feature map respectively, and γ represents the mask rate.

[0063] S32: Element-wise dot product of the down-sampled feature and the designed random mask to obtain the masked feature:

[0064]

[0065] where F(x m ) and F'(x m ) represent the features before and after masking respectively, represents element-wise multiplication.

[0066] In one specific embodiment, the shallow mask reconstruction strategy in Figure 1 is performed on the features obtained from the residual block 2 and the residual block 4 of the residual neural network, where the feature pixel covering rate γ in the mask strategy is set to 15%.

[0067] Step S4: Designing a generator based on convolution and a generator based on perception to reconstruct the masked feature, and using average pooling, linear classification layer and softmax layer to obtain the predicted output of the reconstructed feature, specifically including the following steps:

[0068] S41: Refer to the generator based on convolution in Figure 2 , which mainly contains two 3x3 convolution layers, a GELU activation function layer and a batch normalization layer.

[0069] S42: Refer to the generator based on perception in Figure 2 , which mainly contains two 1x1 convolution layers, a GELU activation function layer and a layer normalization layer.

[0070] S43: Use the generator based on convolution and the generator based on perception to generate features, and the reconstruction method is as follows:

[0071]

[0072] where F(x m ) and F'(x m ) represent the features before and after masking respectively, x m represents the input mixed image data, represents element-wise multiplication.

[0073] S44: Use average pooling operation, linear classifier (also known as linear layer) and multi-classification function softmax to obtain the probability prediction output corresponding to F"(x).

[0074] Step S5: Calculate the difference between the prediction output of the residual neural network and the true label using the cross-entropy function, referred to as the classification loss, and measure the difference between the prediction output of the residual neural network of the current training round and the prediction output of the residual neural network historical model of the last training round using the KL divergence, referred to as the reconstruction loss, and construct a self-distillation loss function by combining the classification loss and the reconstruction loss, specifically including the following steps:

[0075] S51: Calculate the classification loss L ce , the difference between the prediction output of the residual neural network and the true label is calculated using the cross-entropy function, which is calculated as follows:

[0076]

[0077] where K represents the number of classes in the training set, represents the true label y m , p k represents the probability prediction of the kth class of the model.

[0078] S52: Calculate the reconstruction loss L rec1 based on the prediction output of the residual neural network of the current training round, which measures the difference between the prediction output of the reconstruction feature and the prediction output of the residual neural network using the KL divergence, and the calculation process is as follows:

[0079]

[0080] where p k ' represents the probability prediction of the kth class obtained by the reconstruction feature after the average pooling operation, linear classifier and multi-classification function softmax.

[0081] S53: Calculate the reconstruction loss L rec2 based on the residual neural network historical model of the last training round, save the model after each round of training, load the model of the last round when calculating the reconstruction based on the historical model, and use it to predict the corresponding probability prediction and of the mixed images x1 and x2 respectively, and measure the difference between the prediction output of the reconstruction feature and using the KL divergence, and the calculation process is as follows:

[0082]

[0083] S54: The self-distillation loss function constructed contains the classification loss and two reconstruction losses:

[0084] Ltotal =L ce +α(L rec1 +L rec2 )

[0085] Where, L total represents the self-distillation loss function, and α is a hyperparameter used to balance the classification loss and reconstruction loss.

[0086] In a specific embodiment, referring to Figure 1 As shown in the figure, the features of residual block 4 of the residual neural network and the reconstructed features need to be sequentially subjected to an average pooling operation, a linear classifier, and a multi-classification function softmax to obtain the output prediction probability of the residual neural network and the output prediction probability of the reconstructed features. The KL divergence is used to measure the difference between the two. The balance coefficient α between different loss functions is set to 0.5.

[0087] Step S6: Input the enhanced images in the training set into the residual neural network, perform backpropagation training according to the constructed self-distillation loss function, and use the validation set to verify the training effect of the network until the residual neural network converges and saves the optimal network weights.

[0088] In one specific embodiment, the training batch size is set to 128, and images are scaled to 32×32 dimensions and fed into a ResNet18 residual neural network. The training iteration count is set to 200, and the training begins with five epochs of linear warmup with an initial learning rate ranging from 0 to 0.1. The learning rate is multiplied by 0.1 at iterations 105 and 155, respectively, and the SGD optimizer is used. The validation set contains 10,000 images, excluding the training images. Training is terminated when the maximum number of iterations is reached, and the weight file corresponding to the best accuracy and the network weights for each training epoch are saved.

[0089] The above embodiments are preferred implementation modes of the present invention, but the implementation modes of the present invention are not limited to the above embodiments. Any other changes, modifications, substitutions, combinations, and simplifications that do not deviate from the spirit and principles of the present invention should be considered as equivalent replacement methods and are included in the scope of protection of the present invention.

Claims

1. A self-distillation method for image classification based on mask feature reconstruction, characterized in that: The following steps are involved: Step S1: After performing data preprocessing based on image mixing on the input image, the training set and the validation set are divided into two parts according to the proportion, and the training set is enhanced by cropping and rotating. Step S2: Input the preprocessed image into the residual neural network, which contains 4 residual blocks, which are divided into residual block 1, residual block 2, residual block 3 and residual block 4 in sequence. Residual block 1, residual block 2 and residual block 3 are considered to be shallow stages due to the small number of residual structure layers, and residual block 4 is considered to be a deep stage due to the large number of residual structure layers. The features output by residual block 2 and residual block 4 and the predicted output of the residual neural network are obtained through forward propagation, and the features output by residual block 2 are downsampled using convolution with a step size of 4; Step S3: Design a random mask and perform element-wise product of the downsampled features and the random mask to obtain the mask features; Step S4: Use a convolution-based generator and a perceptron-based generator to reconstruct the mask features, and use an average pooling operation, a linear classifier, and a multi-classification function softmax to obtain the predicted output of the reconstructed features; Step S5: Use the cross entropy function to calculate the difference between the predicted output of the residual neural network and the true label, which is called classification loss, and use the KL divergence to measure the difference between the predicted output of the residual neural network in the current training round and the predicted output of the residual neural network historical model in the previous training round, which is called reconstruction loss. By combining the classification loss and the reconstruction loss, a self-distillation loss function is constructed, which includes the following steps: S51: Calculate the classification loss L ce , the cross entropy function is used to calculate the difference between the predicted output of the residual neural network and the true label, which is calculated as follows: In the formula, K represents the number of categories in the training set, Indicates the true label y corresponding to the k-th category m , p k represents the probability prediction of the kth category of the residual neural network; S52: Calculate the reconstruction loss L based on the residual neural network prediction output of the current training round rec1 , the KL divergence is used to measure the difference between the predicted output of the reconstructed features and the predicted output of the residual neural network. The calculation process is as follows: Where p k′ Represents the probability prediction corresponding to the kth category obtained by the reconstructed features after the average pooling operation, linear classifier and multi-classification function softmax; S53: Calculate the reconstruction loss L of the residual neural network history model based on the previous training round rec2 After each round of training, the model is saved. When calculating the reconstruction based on the historical model, the model of the previous round is loaded and used to predict the images x1 and x2 before mixing to obtain the corresponding probability predictions. and KL divergence is used to measure the predicted output of the reconstructed features and The difference between them is calculated as follows: S54: The constructed self-distillation loss function includes both classification loss and two reconstruction losses: L total =L ce +α(L rec1 +L rec2 ) Where, L total represents the self-distillation loss function, α is a hyperparameter used to balance the classification loss and reconstruction loss; Step S6: Input the enhanced images in the training set into the residual neural network, perform backpropagation training according to the constructed self-distillation loss function, and use the validation set to verify the training effect of the network until the residual neural network converges, and save the optimal network weights and the network weights of each training round.

2. The image classification self-distillation method based on mask feature reconstruction according to claim 1, characterized in that: In step S1, two training images x1 and x2 are randomly selected and mixed using linear interpolation: x m =λx1+(1-λ)x2,y m =λy1+(1-λ)y2 Where x m Represents the mixed image after linear interpolation, y1, y2 and y m They represent the true labels corresponding to the two images and the true labels corresponding to the mixed image, and λ is the image mixing ratio.

3. The image classification self-distillation method based on mask feature reconstruction according to claim 2, characterized in that: In step S2, the features output by residual block 2 are downsampled by convolution with a stride of 4 and a convolution kernel of 4 to achieve alignment with the features output by residual block 4 to reduce information loss in the feature map.

4. The image classification self-distillation method based on mask feature reconstruction according to claim 3, characterized in that: The step S3 comprises the following steps: S31: Design of random mask: Where M h,v represents a random mask, R h,v represents a random number between 0 and 1, h and v represent the horizontal and vertical coordinates corresponding to the feature map, and γ represents the mask rate; S32: Perform element-wise dot product of the downsampled features and the designed random mask to obtain the masked features: In the formula, F(x m ) and F′(x m ) represent the features before and after masking, Represents element-wise multiplication.

5. The image classification self-distillation method based on mask feature reconstruction according to claim 4, characterized in that: The step S4 comprises the following steps: S41: Design a convolution-based generator consisting of two 3×3 convolutional layers, a GELU activation function layer, and a batch normalization layer; S42: Design a perceptron-based generator consisting of two 1×1 convolutional layers, a GELU activation function layer, and a layer normalization layer. S43: Use the convolution-based generator and the perceptron-based generator to generate features in sequence. The reconstruction method is as follows: F″(x m )=θ MLP (i conv (F′(x m ))) Where, F″(x m ) represents the reconstructed features, θ conv and θ MLP Represent the convolution-based generator and the perceptron-based generator respectively; S44: Use average pooling operation, linear classifier and multi-classification function softmax to obtain F″(x m ) corresponding probability prediction output.

6. The image classification self-distillation method based on mask feature reconstruction according to claim 5, characterized in that: In step S6, the training batch size is set to 128, the image is scaled to 32×32 size and input into the residual neural network, the number of training iterations is set to 200, and 5 training epochs are started with linear preheating from an initial learning rate of 0 to 0.

1. When the number of iterations reaches 105 and 155, the learning rate is multiplied by 0.1, and the SGD optimizer is used; the validation set is data that does not contain training images, with a total of 10,000 images; when training is performed, the training is stopped when the maximum number of iterations is reached, and the weight file corresponding to the best accuracy and the network weights of each training epoch are saved.

7. The image classification self-distillation method based on mask feature reconstruction according to claim 6, characterized in that: The residual neural network is a ResNet18 residual neural network.

Images