Image classification self-distillation method based on mask feature reconstruction

Abstract

The invention discloses an image classification self-distillation method based on mask feature reconstruction, and the method comprises the steps: carrying out the preprocessing of an input image, dividing the input image into a training set and a verification set, and enhancing the training set; inputting the preprocessed image into a residual neural network, obtaining features output in shallow and deep stages of a residual block and 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; a convolution-based generator and a perceptron-based generator are adopted to reconstruct the mask features, and prediction output of the reconstructed features is obtained; the classification loss and the reconstruction loss are combined to construct a self-distillation loss function; and inputting the image in the enhanced training set into the residual neural network, performing back propagation training according to a self-distillation loss function, verifying by using a verification set, and storing the optimal network weight and the network weight of each training round. According to the method, the network performance is improved, and relatively high training efficiency is kept.

Description

Technical Field

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

[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 technology aims to maintain high accuracy while reducing the number of parameters and computational complexity of high-capacity models. Among them, knowledge distillation is a simple and effective model compression method that forces small models to imitate the knowledge of large models so that small models can achieve the same accuracy as large models. However, the traditional knowledge distillation algorithm is divided into two stages. First, a teacher model needs to be pre-trained, and then the student model is distilled using the teacher 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. In order to solve this problem, the distillation method based on self-learning came into being.

[0003] Although the self-distillation algorithm has shown potential in improving the performance of student models, the existing self-distillation algorithms still have the following major shortcomings: (1) Due to the lack of teacher guidance, the existing self-distillation algorithms often use complex multi-branch complex structures to obtain additional supervision signals, which increases the complexity of the training process to a certain extent and also increases the training time. (2) Most of the existing auxiliary structure-based methods use block stacking to increase the trainability of the shallow layer, ignoring the further optimization of the shallow features. Summary of the invention

[0004] The purpose of the present invention is to overcome the shortcomings and disadvantages of the prior art and provide an image classification self-distillation method based on mask feature reconstruction. The self-distillation strategy is implemented by reconstructing shallow features of the mask, which simplifies the self-distillation process, increases the shallow trainability and optimizes the features, thereby improving the network performance and maintaining a high training efficiency.

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

[0006] Step S1: After performing data preprocessing based on image mixing on the input image, the training set and the validation set are divided according to the proportion, and the training set is subjected to data enhancement by cropping and rotating;

[0007] Step S2: input the preprocessed image into the residual neural network, the residual neural network includes 4 residual blocks, which are respectively 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 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 dot product between the downsampled features and the random mask to obtain the mask features;

[0009] Step S4: using a convolution-based generator and a perceptron-based generator to reconstruct the mask features, and using 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: using the cross entropy function to calculate the difference between the predicted output of the residual neural network and the true label, called classification loss, and using 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, called reconstruction loss. The self-distillation loss function is constructed by combining the classification loss and the reconstruction loss.

[0011] Step S6: Input the enhanced images in the training set into the residual neural network, perform back-propagation 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 by linear interpolation:

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

[0014] In the formula, 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 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] 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;

[0020] S32: Perform element-wise dot product of the downsampled features and the designed random mask to obtain the masked features:

[0021]

[0022] In the formula, F(x m ) and F′(x m ) represent the features before and after masking, respectively. Represents element-wise multiplication.

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

[0024] S41: Design a convolution-based generator consisting of two 3×3 convolutional layers, a GELU activation function layer, and a batch normalization layer;

[0025] S42: Design a perceptron-based generator consisting of two 1×1 convolutional layers, a GELU activation function layer, and a layer normalization layer;

[0026] S43: The convolution-based generator and the perceptron-based generator are used in turn for feature generation, and the reconstruction method is as follows:

[0027] F″(x m )=θ MLP (θ conv (F′(x m )))

[0028] In the formula, F″(x m ) represents the reconstructed features, θ conv and θ MLP Represent the convolution-based generator and the perceptron-based generator respectively;

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

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

[0031] 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:

[0032]

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

[0034] 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:

[0035]

[0036] In the formula, p k ′ represents the probability prediction corresponding to the kth category obtained by reconstructing features through average pooling operation, linear classifier and multi-classification function softmax;

[0037] 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 The KL divergence is used to measure the predicted output of the reconstructed features. The difference between is calculated as follows:

[0038]

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

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

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

[0042] Further, in step S6, the batch size of training 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 of the initial learning rate from 0 to 0.1. When the iterations reach 105 and 155, the learning rate is multiplied by 0.1 respectively, and the SGD optimizer is used; the validation set is data that does not contain training images, which contains 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 weight of each training round are saved.

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

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

[0045] The present invention only designs a single-branch reconstruction auxiliary structure, avoiding the introduction of multi-branch auxiliary modules that complicates training, increasing the trainability of the shallow layer, simplifying the self-distillation process, and improving the training efficiency of the network. In view of the problem that existing algorithms often ignore the optimization of shallow features, the mask feature reconstruction method proposed in the present invention reduces the background information in the shallow features, thereby optimizing the shallow features. BRIEF DESCRIPTION OF THE DRAWINGS

[0046] Figure 1 The network architecture diagram of a specific embodiment of the present invention.

[0047] Figure 2 Architecture diagrams of convolution-based generators and perceptron-based generators. DETAILED DESCRIPTION

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

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

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

[0051] Step S1: After performing data preprocessing based on image mixing on the input image, the training set and the validation set are divided according to the proportion, and the training set is subjected to data enhancement by cropping and rotating;

[0052] In a specific embodiment, CIFAR100 is used as a data set, and a training set and a validation set are divided according to a ratio, wherein the training set contains 50,000 images and the validation set contains 10,000 images.

[0053] Randomly select two training images x1 and x2, and use linear interpolation 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 They represent the true labels corresponding to the two images and the true labels corresponding to the mixed image respectively. λ 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 of the shallow and deep outputs of the residual block and the predicted output of the residual neural network through forward propagation, and downsample the shallow features using convolution with a step size of 4, specifically including the following steps:

[0057] S21: Reference Figure 1 The ResNet18 residual neural network shown in the figure contains 4 residual blocks, which are respectively divided into residual block 1, residual block 2, residual block 3 and residual block 4. 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 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;

[0058] S22: Downsample residual block 2 using convolution with a step size of 4 and a convolution kernel of 4 to achieve alignment with the output features of residual block 4. The convolution-based downsampling method can effectively reduce information loss in the feature map, and compared with the pooling-based downsampling method, it can reduce 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-wise dot product between the downsampled features and the random mask to obtain mask features, specifically including 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: Perform element-wise dot product of the downsampled features and the designed random mask to obtain the masked features:

[0064]

[0065] In the formula, F(x m ) and F′(x m ) represent the features before and after masking, respectively. Represents element-wise multiplication.

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

[0067] Step S4: Design a convolution-based generator and a perceptron-based generator to reconstruct the mask features, and use average pooling, linear classification layer and softmax layer to obtain the prediction output of the reconstructed features, which specifically includes the following steps:

[0068] S41: Reference Figure 2 The convolution-based generator in mainly consists of two 3×3 convolutional layers, a GELU activation function layer, and a batch normalization layer.

[0069] S42: Reference Figure 2 The perceptron-based generator in Figure 1 mainly consists of two 1×1 convolutional layers, a GELU activation function layer, and a layer normalization layer.

[0070] S43: Feature generation using a convolution-based generator and a perceptron-based generator, the reconstruction is as follows:

[0071]

[0072] In the formula, 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 the average pooling operation, the linear classifier (also called the linear layer) and the multi-classification function softmax to obtain the probability prediction output corresponding to F″(x).

[0074] Step S5: 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 called the classification loss, and the KL divergence is used 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 the reconstruction loss. The self-distillation loss function is constructed by combining the classification loss and the reconstruction loss, which specifically includes the following steps:

[0075] 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:

[0076]

[0077] In the formula, K represents the number of categories in the training set, Represents the true label y corresponding to the kth category m , p k represents the probability prediction of the kth category of the model.

[0078] 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:

[0079]

[0080] In the formula, p k ′ represents the probability prediction corresponding to the kth category obtained by reconstructing the features through average pooling operation, linear classifier and multi-classification function softmax.

[0081] 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 The KL divergence is used to measure the predicted output of the reconstructed features. The difference between is calculated as follows:

[0082]

[0083] S54: The constructed self-distillation loss function includes both 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 the 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 go through an average pooling operation, a linear classifier, and a multi-classification function softmax in turn to obtain the output prediction probability of the residual neural network and the output prediction probability of the reconstructed features, and 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 back-propagation 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 a specific embodiment, the batch size of training is set to 128, the image is scaled to 32×32 size and input into the ResNet18 residual neural network, the number of training iterations is set to 200, and 5 training rounds of epochs are started with linear preheating of the initial learning rate from 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, and contains 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 weight of each training round 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 equivalent replacement methods and are included in the protection scope 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 according to the proportion, and the training set is subjected to data enhancement by cropping and rotating; Step S2: input the preprocessed image into the residual neural network, the residual neural network includes 4 residual blocks, which are respectively 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 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 dot product between the downsampled features and the random mask to obtain the mask features; Step S4: using a convolution-based generator and a perceptron-based generator to reconstruct the mask features, and using an average pooling operation, a linear classifier, and a multi-classification function softmax to obtain the predicted output of the reconstructed features; Step S5: using the cross entropy function to calculate the difference between the predicted output of the residual neural network and the true label, called classification loss, and using 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, called reconstruction loss. The self-distillation loss function is constructed by combining the classification loss and the reconstruction loss. Step S6: Input the enhanced images in the training set into the residual neural network, perform back-propagation 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 In the formula, 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 step size 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, respectively, 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, respectively. 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: The convolution-based generator and the perceptron-based generator are used in turn for feature generation, and the reconstruction method is as follows: F″(x m )=θ MLP (i conv (F′(x m ))) In the formula, 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: The step S5 comprises 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, Represents the true label y corresponding to the kth 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: In the formula, p k′ It represents the probability prediction corresponding to the kth category obtained by reconstructing features through 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 The KL divergence is used to measure the predicted output of the reconstructed features. The difference between 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, and α is a hyperparameter used to balance the classification loss and the reconstruction loss.

7. The image classification self-distillation method based on mask feature reconstruction according to claim 6, characterized in that: In step S6, the batch size of training 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 of the initial learning rate from 0 to 0.

1. When the number of iterations reaches 105 and 155, the learning rate is multiplied by 0.1 respectively, and the SGD optimizer is used; the validation set is data that does not contain training images, which contains 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 weight of each training round are saved.

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

Images