A self-distillation image classification method based on generative adversarial networks and mutual learning

Abstract

This invention discloses a self-distillation image classification method based on generative adversarial networks (GANs) and mutual learning. To make the features extracted by the shallow student network approximate those extracted by the deep network, we utilize GANs. The shallow student network, due to its simple structure and weak feature extraction ability, has its extracted features defined as generated samples. The backbone network, with its complex structure and deeper layers, has a stronger feature extraction ability; its extracted features are defined as real samples. Based on this, we introduce a pre-designed discriminator into the self-distillation framework, allowing the generator (shallow student network) to compete with the discriminator, thereby improving the feature extraction ability of the student network. To ensure stable training of the GAN, we use student features and teacher features before the activation layer to train the discriminator. Furthermore, in the self-distillation framework, the student and teacher networks are trained from random initialization. The learning performance of the student network cannot be fed back to the teacher network, making it difficult for the teacher network to optimize and adjust its knowledge. Therefore, we introduce a mutual learning training method.

Description

Technical Field

[0001] This invention relates to the fields of neural network model compression, generative adversarial networks, mutual learning, and image classification. In particular, it relates to a self-distillation image classification method based on generative adversarial networks and mutual learning. Background Technology

[0002] Deep learning has developed rapidly in recent years and has been widely applied in tasks such as image recognition and object detection. Convolutional neural networks (CNNs), as a crucial component of deep learning, typically consist of multiple convolutional layers. The more layers a CNN has, the more parameters it generates. While increasing the number of layers improves performance, it also increases computational costs, requiring high computing power for both training and deployment. This limits the network's applicability on resource-constrained devices. Model compression aims to reduce the number of parameters and computational cost of neural networks, reducing model complexity without altering its original performance, thus enabling its application on edge devices.

[0003] The innovation of this invention and the difference from the prior art

[0004] In the self-distillation framework, to enable shallow student networks to achieve the same feature extraction capabilities as deep teacher networks, a self-distillation method integrating generative adversarial networks (GANs) and mutual learning is proposed. First, features extracted by the student network are treated as generated data, while features extracted by the teacher network are treated as real data. Then, an additional discriminator is introduced into the self-distillation framework. Through continuous interaction between the generator and the discriminator, the feature similarity between the student and teacher networks is improved. Second, the discriminator is input with the student and teacher features before the activation layers, ensuring the integrity of feature information during training. Finally, since the parameters of the student and teacher networks are randomly initialized before training in the self-distillation framework, and knowledge transfer is a unidirectional process (only the teacher transfers knowledge to the student), the learning quality of the student network cannot be fed back to the teacher network, making it difficult for the teacher network to optimize and adjust its knowledge. Therefore, a mutual learning mechanism is introduced. This method can effectively improve the classification performance of shallow networks. Summary of the Invention

[0005] Existing deep learning models have a large number of parameters and are difficult to deploy. To address this issue, this invention employs self-distillation, combining generative adversarial networks (GANs) and mutual learning methods. This method enhances the feature extraction capability of shallow branches. In application, prediction can be performed using only shallow branches, or deep and shallow branches can be combined for joint prediction regardless of the number of parameters.

[0006] A self-distillation image classification method based on generative adversarial networks and mutual learning mainly includes the following steps:

[0007] Step S1: Preprocess the training data

[0008] Step S1.1: Augment the training data using random pruning and horizontal flipping.

[0009] Step S1.2: Randomly shuffle the data and divide it into different batches.

[0010] Step S2, Model Training:

[0011] Step S2.1: Introduce shallow branch classification exits in the neural network. The shallow branch acts as the student network, and the deep branch classification exit acts as the teacher network. The shallow and deep branches constitute the general structure of self-distillation.

[0012] Step S2.2: Design a discriminator, which consists of multiple convolutional layers and a final fully connected layer. Features extracted by shallow branches of the neural network are considered fake samples, while features extracted by deep branches are considered real samples.

[0013] Step S2.3: Input the features extracted from the shallow and deep branches into the discriminator respectively, and use the idea of ​​generative adversarial network to make the two compete with each other. The goal is to make the features extracted from the shallow branch close to the features extracted from the deep branch.

[0014] Step S2.4: When the discriminator determines that the feature of the shallow branch is true, the loss suffered by the discriminator due to the incorrect judgment is defined as loss. D At this point, the features extracted by the shallow branch are already sufficiently good. When the discriminator determines that the features of the shallow branch are false, the shallow branch incurs a loss defined as [loss value missing]. S1 .

[0015] Step S2.4 introduces a mutual learning training method on both the shallow and deep branches, constructing a mutual learning loss to allow information exchange between the two branches. The mutual learning loss received by the shallow branch is loss2, and the mutual learning loss received by the deep branch is loss3.

[0016] In step S2.5, the features extracted from the shallow branch are fed into its classifier to obtain classification result o1, and the features extracted from the deep branch are fed into its classifier to obtain classification result o2. The two classification results are then processed by softmax to obtain the predicted values ​​for the samples. The cross-entropy losses between the predicted values ​​and the true labels are loss4 and loss5, respectively.

[0017] In step S2.6, during training, the learning rate is decayed according to different rounds. The initial learning rate is set to 0.1, and at rounds 125, 260, and 420, the learning rate becomes one-tenth of the original. L2 regularization is also introduced. A multi-module loss is used, with the neural network subjected to generative adversarial loss, mutual learning loss, and label loss.

[0018] Step S2.7, forward propagation. After forward propagation is completed, the total loss value for this batch can be calculated.

[0019] Step S2.8, Backpropagation. The neural network is trained using the stochastic gradient descent (SGD) method. One round ends after all batches have completed one backpropagation pass.

[0020] Step S2.8: Repeat the training process described above until the model finally converges.

[0021] Step S3, Prediction:

[0022] Step S3.1: Remove the discriminator from the neural network. Prune the network, keeping only the shallow branches that predict the exit.

[0023] Step S3.2: Input the image to be predicted into the model to obtain the classification result. Attached Figure Description

[0024] Figure 1 This is a structural diagram of the model involved in this invention.

[0025] Figure 2 This is a flowchart of the actual model processing used in this invention. Detailed Implementation

[0026] To make the objectives, technical solutions and advantages of the present invention clearer, the present invention will be further described in detail below with reference to specific embodiments and accompanying drawings.

[0027] This embodiment uses the Cifar100 dataset as the training dataset for image classification. The Cifar100 dataset contains 60,000 training images, of which 50,000 are the training set and 10,000 are the test set, with a total of 10 categories.

[0028] Step S1: Preprocess the Cifar100 training set.

[0029] Step S1.1: Augment the training data using random pruning and horizontal flipping.

[0030] Step S1.2: Randomly shuffle the training set data and divide it into different batches.

[0031] Step S2, Model Training:

[0032] Step S2.1 introduces a shallow branch classification exit in the ResNet18 neural network. The shallow branch acts as the student network, and the deep branch classification exit acts as the teacher network. The shallow and deep branches constitute the general structure of self-distillation.

[0033] Step S2.2: Design the discriminator, which consists of multiple convolutional layers and a final fully connected layer. Its purpose is to distinguish between real and fake samples, classifying them as either true or false, represented by labels 1 or 0. Features extracted by shallow branches of the neural network are considered fake samples, while features extracted by deep branches are considered real samples. In ResNet18, the dimensions of real and fake samples are 512×1×1.

[0034] Step S2.3: Using the idea of ​​generative adversarial networks, the two branches engage in a game. The features extracted by the deep and shallow branches are input into the discriminator, respectively. The purpose is to make the features extracted by the shallow branch closer to the features extracted by the deep branch by judging the ownership of the sample.

[0035] Step S2.4: When the discriminator determines that the feature of the shallow branch is true, the loss suffered by the discriminator due to the incorrect judgment is defined as loss. D At this point, the features extracted by the shallow branch are already sufficiently good. When the discriminator determines that the features of the shallow branch are false, the shallow branch incurs a loss defined as [loss value missing]. S1 .

[0036] Step S2.4 introduces a mutual learning training method on both the shallow and deep branches, constructing a mutual learning loss to allow information exchange between the two branches. The mutual learning loss received by the shallow branch is loss2, and the mutual learning loss received by the deep branch is loss3.

[0037] In step S2.5, the features extracted from the shallow branch are fed into its classifier to obtain classification result o1, and the features extracted from the deep branch are fed into its classifier to obtain classification result o2. The two classification results are then processed by softmax to obtain the predicted values ​​for the samples. The cross-entropy losses between the predicted values ​​and the true labels are loss4 and loss5, respectively.

[0038] In step S2.6, during training, the learning rate is decayed according to different rounds. The initial learning rate is set to 0.1, and at rounds 125, 260, and 420, the learning rate becomes one-tenth of the original. L2 regularization is also introduced. A multi-module loss is used, with the neural network subjected to generative adversarial loss, mutual learning loss, and label loss.

[0039] Step S2.7, forward propagation. After forward propagation is completed, the total loss value for this batch can be calculated.

[0040] Step S2.8, Backpropagation. The neural network is trained using the stochastic gradient descent (SGD) method. One round ends after all batches have completed one backpropagation pass.

[0041] Step S2.8: Repeat the training process described above until the model finally converges.

[0042] Step S3, Prediction:

[0043] Step S3.1: Remove the discriminator from the neural network. In the prediction phase, only the forward propagation operation is needed to remove the loss function and the discriminator. When resources are scarce, the network can be pruned, retaining only the shallow branch to predict the output. When resources are sufficient, both the shallow and deep branches can be used for joint prediction.

[0044] Step S3.2: Input the predicted data of the dataset into the neural network to obtain the classification result.

[0045] This paper proposes a self-distillation image classification method based on generative adversarial networks (GANs) and mutual learning. To make the features extracted by the shallow student network approximate those extracted by the deep network, we utilize GANs. The shallow student network, due to its simple structure and weak feature extraction ability, has its extracted features defined as generated samples. The backbone network, with its complex structure and deeper layers, has a stronger feature extraction ability; its extracted features are defined as real samples. Based on this, we introduce a pre-designed discriminator into the self-distillation framework, allowing the generator (shallow student network) to compete with the discriminator, thereby improving the feature extraction ability of the student network. To ensure stable training of the GAN, we use student features and teacher features before the activation layer to train the discriminator. Furthermore, since the student and teacher networks in the self-distillation framework are trained from random initialization, the learning performance of the student network cannot be fed back to the teacher network, making it difficult for the teacher network to optimize and adjust its knowledge. Therefore, we introduce a mutual learning training method.

[0046] The above description is merely a specific embodiment of the present invention, but the scope of protection of the present invention is not limited thereto. Any transformations or substitutions that can be conceived by those skilled in the art within the scope of the technology disclosed in the present invention should be included within the scope of the present invention. Therefore, the scope of the present invention should be determined by the scope of protection of the patent document.

Claims

1. A self-distillation image classification method based on generative adversarial networks and mutual learning, characterized in that... The three processes—data preprocessing, model training, and prediction—are detailed below: Step S1: Preprocess the training data Step S1.1: Augment the training data using random pruning and horizontal reversal. Step S1.2: Randomly shuffle the data and divide it into different batches; Step S2, Model Training: Step S2.1: Introduce shallow branch classification exits in the neural network; shallow branches are treated as student networks, and deep branch classification exits are treated as teacher networks; shallow and deep branches constitute the general structure of self-distillation. Step S2.2: Design a discriminator, which consists of multiple convolutional layers and a final fully connected layer; treat features extracted by shallow branches of the neural network as fake samples and features extracted by deep branches of the neural network as real samples. Step S2.3: Input the features extracted from the shallow and deep branches into the discriminator respectively, and use the idea of ​​generative adversarial network to make the two compete with each other. The goal is to make the features extracted from the shallow branch close to the features extracted from the deep branch. Step S2.4: When the discriminator determines that the feature of the shallow branch is true, the loss suffered by the discriminator due to the incorrect judgment is defined as loss. D At this point, the features extracted by the shallow branches are already good enough. When the discriminator determines that the feature of the shallow branch is false, the loss incurred by the shallow branch is defined as loss. S1 ; Step S2.4: Introduce a mutual learning training method on the shallow and deep branches, and construct a mutual learning loss to enable the two branches to exchange information; the mutual learning loss received by the shallow branch is loss2, and the mutual learning loss received by the deep branch is loss3. Step S2.5: The features extracted from the shallow branch are fed into its classifier to obtain classification result o1, and the features extracted from the deep branch are fed into its classifier to obtain classification result o2; the two classification results are processed by softmax to obtain the sample prediction value; the cross-entropy loss between the prediction value and the true label is loss4 and loss5 respectively. Step S2.6: During training, the learning rate is decayed according to different rounds. The initial learning rate is set to 0.1, and the learning rate becomes one-tenth of the original rate at rounds 125, 260, and 420. At the same time, L2 regularization is introduced. The neural network is subjected to generative adversarial loss, mutual learning loss, and label loss in combination with multi-module loss. Step S2.7, forward propagation; after forward propagation is completed, the total loss value of this batch can be obtained; Step S2.8, backpropagation; use the stochastic gradient descent (SGD) method to train the neural network; once all batches have completed one backpropagation, one round ends; Step S2.8: Repeat the training process described above until the model finally converges; Step S3, Prediction: Step S3.1: Remove the discriminator from the neural network; prune the network, keeping only the shallow branches that predict the exit; Step S3.2: Input the image to be predicted into the model to obtain the classification result.

Images