ViT distillation training method, system and device and readable storage medium

Abstract

The invention provides a ViT distillation training method, system and device and a readable storage medium, and the method comprises the steps: selecting a long-tail data set, and carrying out the preprocessing of the long-tail data set; training the convolutional neural network to generate a teacher model; generating a data sample based on the long-tail data set, inputting the data sample into a teacher model, and generating a soft label; constructing a student model based on a visual Transform, and configuring a knowledge distillation function and a loss function for the student model; generating a data sample based on the long-tail data set, inputting the data sample into a student model, generating a classification result, and optimizing student model parameters by using a loss function based on the classification result and the soft label to train the student model; and after training is completed, performing model verification and output on the student model. According to the method, the technologies of long-tail data set preprocessing, teacher network training, knowledge distillation, loss reweighting and the like are combined, so that the classification performance and generalization ability of the ViT model on the long-tail data set are effectively improved.

Description

Technical Field

[0001] The present invention relates to the field of computer vision and deep learning technology, and more specifically to a ViT distillation training method, system, device and readable storage medium. Background Art

[0002] With the rapid development of deep learning technology, Vision Transformer (ViT), as an emerging deep learning model, has shown strong performance in image recognition, classification and other tasks. However, in practical applications, especially when dealing with long-tailed datasets, the ViT model faces many challenges. Long-tailed datasets refer to datasets in which the number of samples of certain categories is much larger than that of other categories. This unbalanced data distribution often leads to a significant decrease in the performance of the model on a minority of categories, which is the so-called "long-tail effect".

[0003] Traditional deep learning models, including Convolutional Neural Networks (CNNs), also have similar problems on long-tail datasets. Although CNNs have made remarkable achievements in image recognition, their reliance on local features makes the model prone to overfitting when dealing with minority categories, especially when the number of samples in these categories is limited.

[0004] To overcome this challenge, researchers have proposed a variety of methods, such as resampling, reweighted loss functions, etc., to balance the contributions of different categories during training. However, while these methods improve the model's ability to recognize minority categories, they may also have a negative impact on the performance of majority categories.

[0005] Knowledge distillation, as a model compression and transfer learning technology, has received widespread attention in the field of deep learning in recent years. By transferring the knowledge of a complex teacher model to a simple student model, knowledge distillation can reduce the complexity and computational cost of the model while maintaining or improving model performance. In the context of long-tail datasets, knowledge distillation technology is expected to help the student model improve its recognition ability for minority categories through the effective learning of minority categories by the teacher model.

[0006] However, most existing knowledge distillation methods are designed for balanced datasets and have limited effects when directly applied to long-tail datasets. Therefore, a knowledge distillation training method specifically optimized for long-tail datasets is needed to effectively improve the classification performance and generalization ability of the ViT model on long-tail datasets. Summary of the invention

[0007] In view of the above problems, the purpose of the present invention is to provide a ViT distillation training method, system, device and readable storage medium, which effectively improves the classification performance and generalization ability of the ViT model on long-tail datasets by combining long-tail dataset preprocessing, teacher network training, knowledge distillation and loss reweighting techniques.

[0008] In order to achieve the above object, the present invention is implemented through the following technical solutions:

[0009] In a first aspect, the present invention discloses a ViT distillation training method, comprising:

[0010] Select a long-tail data set and preprocess the long-tail data set;

[0011] Train the convolutional neural network to generate a teacher model;

[0012] Generate data samples based on the long-tail dataset, input them into the teacher model, and generate soft labels;

[0013] Build a student model based on the visual Transformer and configure the knowledge distillation function and loss function for the student model;

[0014] Generate data samples based on the long-tail data set, input them into the student model, generate classification results, and optimize the student model parameters using the loss function based on the classification results and soft labels to train the student model;

[0015] After the training is completed, the student model is verified and output.

[0016] Further, the selecting of the long-tail data set and preprocessing of the long-tail data set include:

[0017] Choose CIFAR-10-LT or ImageNet-LT as the long-tail dataset;

[0018] Oversample the tail categories of the long-tail dataset and undersample the head categories of the long-tail dataset to generate data samples;

[0019] Adjust the class weights according to the number of data samples in each class and reweight the loss of the tail classes when calculating the loss function;

[0020] All sample data are weakly enhanced, and sample data of tail categories are strongly enhanced to generate out-of-distribution images.

[0021] Further, the training of the convolutional neural network to generate a teacher model includes:

[0022] Train the convolutional neural network CSPDarkNet;

[0023] During training, the model parameters are updated by calculating the SAM perturbation direction and the gradient after the perturbation;

[0024] After training is completed, a teacher model is generated.

[0025] Further, generating data samples based on the long-tail data set, inputting into the teacher model, and generating soft labels include:

[0026] The data samples after strong enhancement are input into the teacher model to generate soft labels.

[0027] Furthermore, the student model is constructed based on the visual Transformer, and a knowledge distillation function and a loss function are configured for the student model, including:

[0028] Use visual Transformer as the student model;

[0029] DS Token is introduced into the student model for distillation learning, and DS Token is used to learn features by imitating the soft labels of the teacher model;

[0030] Configure the loss function for calculating classification loss and distillation loss for the student model;

[0031] The classification loss is calculated using cross entropy loss, and the distillation loss is calculated by minimizing the KL divergence function.

[0032] Further, the method of optimizing the student model parameters using a loss function based on the classification results and the soft labels includes:

[0033] Based on the classification results and soft labels, the classification loss and distillation loss are calculated using the loss function to jointly optimize the parameters of the student model;

[0034] When calculating the distillation loss, the weight of the soft labels of the tail category sample data is increased to ensure the attention of the tail category sample data.

[0035] Further, the weak enhancement processing includes:

[0036] Random cropping processing, horizontal flipping processing and color adjustment processing; the strong enhancement processing includes: random rotation processing, cropping processing with a preset amplitude and color dithering processing.

[0037] In a second aspect, the present invention further discloses a ViT distillation training system, comprising:

[0038] A preprocessing module is used to select and preprocess the long-tail data set;

[0039] The teacher model building module is used to train the convolutional neural network and generate the teacher model;

[0040] The soft label generation module is used to generate data samples based on the long-tail data set, input the teacher model, and generate soft labels;

[0041] The student model building module is used to build a student model based on the visual Transformer and configure the knowledge distillation function and loss function for the student model;

[0042] The model training module is used to generate data samples based on the long-tail data set, input the data into the student model, generate classification results, and optimize the student model parameters using the loss function based on the classification results and soft labels to train the student model;

[0043] The model verification output module is used to verify and output the student model after training is completed.

[0044] In a third aspect, the present invention further discloses a ViT distillation training device, comprising:

[0045] A memory for storing the ViT distillation training program;

[0046] A processor is used to implement the steps of the ViT distillation training method as described in any one of the above items when executing the ViT distillation training program.

[0047] In a fourth aspect, the present invention further discloses a readable storage medium, on which a ViT distillation training program is stored. When the ViT distillation training program is executed by a processor, the steps of the ViT distillation training method described in any one of the above items are implemented.

[0048] Compared with the prior art, the present invention has the following beneficial effects:

[0049] 1. The present invention introduces knowledge distillation technology to pass the soft labels generated by the teacher model (CNN) to the student model (ViT), thereby enhancing the student model's ability to extract local features, especially the feature learning ability for minority categories.

[0050] 2. The present invention introduces DS Token as a distillation regularization tool in the student model, which improves the effect of ViT in the distillation process and further improves the adaptability of the model to long-tail data sets.

[0051] 3. The present invention uses SAM technology to optimize the teacher model. By finding a flat area in the parameter space for training, the overfitting of the model to the training data is reduced, and the robustness and generalization ability of the model on long-tail data sets are improved.

[0052] 4. The present invention discloses an optimization strategy specifically for training ViT in long-tail data sets, using CNN as the teacher network and ViT as the student network for distillation training; by using CSPDarkNet as the CNN teacher network to guide the student network training optimization, by obtaining a flat optimal solution through low-rank feature learning when training CNN, the CNN network's resistance to redundant features and generalization ability are improved. When training ViT, in addition to CLS token and token as input, a DS token is added as a feature extraction optimizer for distillation training. The final loss function is the weighted sum of ViT classification loss and distillation loss. Therefore, ViT can fully exert its ability to capture long-range features, and can also learn CNN's ability to capture local features, enhancing the model's classification ability for long-tail data sets.

[0053] It can be seen that compared with the prior art, the present invention has outstanding substantive features and significant progress, and the beneficial effects of its implementation are also obvious. BRIEF DESCRIPTION OF THE DRAWINGS

[0054] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings described below are only embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on the provided drawings without paying creative work.

[0055] Figure 1 It is a method flow chart of a specific implementation mode of the present invention.

[0056] Figure 2 It is a system structure diagram of a specific implementation mode of the present invention.

[0057] Figure 3 It is a schematic diagram of the device structure of a specific embodiment of the present invention.

[0058] In the figure, 1. preprocessing module; 2. teacher model building module; 3. soft label generation module; 4. student model building module; 5. model training module; 6. model verification output module; 101. processor; 102. memory; 103. input interface; 104. output interface; 105. communication unit; 106. keyboard; 107. display; 108. mouse. DETAILED DESCRIPTION

[0059] Unless otherwise defined, all technical and scientific terms used herein have the same meaning as those commonly understood by those skilled in the art of the present invention. The terms used in the specification of the present invention herein are only for the purpose of describing specific embodiments and are not intended to limit the present invention.

[0060] The key terms appearing in the present invention are explained below.

[0061] Long-Tailed Datasets: Long-Tailed Datasets refer to datasets with extremely unbalanced category distribution. The main feature of long-tailed datasets is the imbalanced category distribution. Some categories contain a large number of samples (called head categories), while other categories contain very few samples (called tail categories). The number of samples in the dataset decreases rapidly as the number of categories increases, forming a long-tail distribution. The number of samples in the head category is much larger than that in the tail category. The small number of samples in the tail category makes it difficult for the model to learn the characteristics of these categories, thereby reducing the classification accuracy of the model for these categories.

[0062] Distillation: is a model compression and transfer learning technique that extracts knowledge from a large, complex pre-trained model (called the teacher model) and transfers that knowledge to a smaller, simpler model (called the student model).

[0063] Vision Transformer: Vision Transformer (ViT) is an image classification model based on the Transformer architecture.

[0064] DS token: Like the ViT model input vector patch tokens and CLS token, as the input vector of the transformer input, the initial value of DSTToken is a learnable vector, which is randomly initialized.

[0065] In order to enable those skilled in the art to better understand the scheme of the present invention, the present invention is further described in detail below in conjunction with the accompanying drawings and specific implementation methods. Obviously, the described embodiments are only part of the embodiments of the present invention, rather than all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative work are within the scope of protection of the present invention.

[0066] See also Figure 1 As shown, this embodiment provides a ViT distillation training method, comprising the following steps:

[0067] S1: Select a long-tail dataset and preprocess it.

[0068] In a specific implementation, first, CIFAR-10-LT or ImageNet-LT is selected as a long-tail dataset. At this time, the long-tail dataset is preprocessed, including: oversampling the tail categories of the long-tail dataset and undersampling the head categories of the long-tail dataset to generate data samples; adjusting the category weights according to the number of data samples of each category and reweighting the loss of the tail categories when calculating the loss function; weakly enhancing all sample data and strongly enhancing the sample data of the tail categories to generate out-of-distribution images.

[0069] As an example, first select a suitable long-tail dataset, such as CIFAR-10-LT, ImageNet-LT, etc. Make sure that the dataset contains a variety of categories and the category distribution is unbalanced.

[0070] Then, the long-tail datasets are preprocessed by resampling, reweighting and data augmentation to improve the performance of Visual Transformer (ViT) on these datasets.

[0071] The preprocessing of long-tail data sets is mainly carried out in the following ways:

[0072] ①Resampling:

[0073] Oversampling: Oversampling the tail classes (classes with fewer samples) to increase the number of samples in these classes, thereby balancing the class distribution.

[0074] Undersampling: Undersampling the head categories (categories with more samples) to reduce the number of samples in these categories and thus balance the category distribution.

[0075] Among them, the ratio of oversampled tail categories and undersampled head categories is dynamically adjusted according to the imbalance of category distribution of the dataset.

[0076] ②Re-weighting:

[0077] Class weight adjustment: Based on the number of samples in a class, higher weights are assigned to tail classes to ensure that the model pays more attention to these classes during training. Class weight adjustment can use the inverse class frequency as a weight factor to further focus on tail classes.

[0078] Loss function reweighting: While calculating the loss function, higher weights are given to the losses of the tail classes, thereby increasing their contribution in the total loss.

[0079] ③Data Augmentation:

[0080] Weak Augmentation: Apply basic data augmentation techniques such as random cropping, horizontal flipping, and color adjustment to all data samples.

[0081] Strong Augmentation: Especially for the data samples in the tail category, stronger data augmentation techniques are used, such as random rotation, large-scale cropping, and color jittering, to generate out-of-distribution (OOD) images.

[0082] S2: Train the convolutional neural network to generate a teacher model.

[0083] In a specific implementation, the convolutional neural network CSPDarkNet is trained; during training, the model parameters are updated by calculating the SAM perturbation direction and the gradient after the perturbation; after the training is completed, a teacher model is generated.

[0084] As an example, a pre-trained convolutional neural network (CNN) is used as a teacher model, and the present invention uses the classic backbone CSPDarkNet of the YOLO series algorithm as a teacher model. The parameters are initialized using the pre-trained model, which helps to speed up the training process and improve the initial performance.

[0085] When training the network, the model parameters are updated in the following two ways:

[0086] Calculation of SAM perturbation direction: Calculate the current gradient and find a perturbation direction that increases the loss the most in the parameter space. Move a certain distance in this direction to calculate the new loss.

[0087] Gradient calculation after perturbation: After moving a certain distance in the found perturbation direction, calculate the new loss and use this loss to calculate the gradient after perturbation. Use the gradient after perturbation to update the model parameters to find a flatter optimal solution in the parameter space.

[0088] S3: Generate data samples based on the long-tail dataset, input them into the teacher model, and generate soft labels.

[0089] In a specific implementation, the data samples after strong enhancement processing are input into the teacher model to generate soft labels (ie, probability distribution).

[0090] S4: Build a student model based on the visual Transformer and configure the knowledge distillation function and loss function for the student model.

[0091] In a specific implementation, a visual Transformer is used as a student model. DS Token for distillation learning is introduced into the student model, and features are learned by imitating the soft labels of the teacher model using DS Token. At the same time, a loss function for calculating classification loss and distillation loss is configured for the student model; specifically, the classification loss is calculated using cross entropy loss, and the distillation loss is calculated using the minimized KL divergence function.

[0092] As an example, this step specifically includes the following methods:

[0093] Student model architecture: Use Visual Transformer (ViT) as the student model.

[0094] Introducing DS Token: Introduce a DS Token specifically for distillation learning in the student model. In addition to the standard classification Token (CLS Token), this Token learns features by imitating the soft labels of the teacher model.

[0095] Loss function design: During the training process, the loss function of the student model consists of two parts. The classification loss uses the cross entropy loss. The distillation loss is calculated by minimizing the Kullback-Leibler divergence (KL divergence) using the difference between the soft labels generated by the teacher model.

[0096] S5: Generate data samples based on the long-tail data set, input them into the student model, generate classification results, and use the loss function to optimize the student model parameters based on the classification results and soft labels to train the student model.

[0097] In a specific implementation, when using the loss function for model training, the loss function is used to calculate the classification loss and the distillation loss based on the classification results and the soft labels to jointly optimize the parameters of the student model.

[0098] Specifically, when calculating the distillation loss, a higher weight is given to the soft labels of the tail categories to ensure that these categories receive more attention during training. During training, a task separation strategy is adopted to separate the tasks of classification tokens (CLS tokens) and distillation tokens (DS tokens). Classification tokens focus on the classification task of the overall category, while DS tokens focus on the knowledge in the soft labels learned from the teacher model, especially the feature representation for the tail categories.

[0099] S6: After training is completed, the student model is verified and output.

[0100] In a specific implementation, a validation set is generated based on a long-tail dataset, and the performance of the student model is evaluated on the validation set to ensure its generalization ability on the long-tail dataset. After the training is completed, the final student model is saved and used for object detection or classification tasks in practical applications. Among them, the classification accuracy, recall rate and F1 score of the student model on the long-tail dataset can be evaluated on the validation set to comprehensively evaluate the performance of the model.

[0101] As an example, save the parameters and configuration files of the student model after training. The trained student model can be deployed to actual application scenarios for computer vision tasks such as object detection and image classification, especially for applications on long-tail datasets.

[0102] The present invention provides a ViT distillation training method, which effectively improves the classification performance and generalization ability of the ViT model on long-tail datasets by combining long-tail dataset preprocessing, teacher network training, knowledge distillation, and loss reweighting techniques.

[0103] See also Figure 2 As shown, the present invention also discloses a ViT distillation training system, including: a preprocessing module 1, a teacher model construction module 2, a soft label generation module 3, a student model construction module 4, a model training module 5 and a model verification output module 6.

[0104] The preprocessing module 1 is used to select a long-tail data set and preprocess the long-tail data set.

[0105] The teacher model building module 2 is used to train the convolutional neural network and generate a teacher model.

[0106] The soft label generation module 3 is used to generate data samples based on the long-tail data set, input the teacher model, and generate soft labels.

[0107] The student model building module 4 is used to build a student model based on the visual Transformer and configure the knowledge distillation function and loss function for the student model.

[0108] The model training module 5 is used to generate data samples based on the long-tail data set, input the data into the student model, generate classification results, and optimize the student model parameters using the loss function based on the classification results and soft labels to train the student model.

[0109] The model verification output module 6 is used to verify and output the student model after the training is completed.

[0110] The specific implementation of the ViT distillation training system of this embodiment is basically the same as the specific implementation of the ViT distillation training method described above, and will not be repeated here.

[0111] See also Figure 3As shown, the present invention also discloses a ViT distillation training device, comprising a processor 101 and a memory 102; wherein, when the processor 101 executes the ViT distillation training program stored in the memory, the steps of the ViT distillation training method as described in any one of the above items are implemented.

[0112] The ViT distillation training device provided in this embodiment may include but is not limited to a smart phone, a tablet computer, a laptop computer or a desktop computer.

[0113] Among them, the processor 101 may include one or more processing cores, such as a 4-core processor, an 8-core processor, etc. The processor 101 can be implemented in at least one hardware form of a digital signal processor (DSP), a field-programmable gate array (FPGA), and a programmable logic array (PLA). The processor 101 may also include a main processor and a coprocessor. The main processor is a processor for processing data in an awake state, also known as a central processing unit (CPU); the coprocessor is a low-power processor for processing data in a standby state. In some embodiments, the processor 101 may be integrated with a graphics processing unit (GPU), and the GPU is responsible for rendering and drawing the content to be displayed on the display screen. In some embodiments, the processor 101 may also include an artificial intelligence (AI) processor, which is used to process computing operations related to machine learning.

[0114] The memory 102 may include one or more computer-readable storage media, which may be non-transitory. The memory 102 may also include a high-speed random access memory, and a non-volatile memory, such as one or more disk storage devices, flash memory storage devices. In this embodiment, the memory 102 is at least used to store the following computer program, wherein, after the computer program is loaded and executed by the processor 101, it can implement the relevant steps of the ViT distillation training method disclosed in any of the aforementioned embodiments. In addition, the resources stored in the memory 102 may also include an operating system and data, etc., and the storage method may be short-term storage or permanent storage. Among them, the operating system may include Windows, Unix, Linux, etc. The data may include, but is not limited to, the data involved in the above-mentioned ViT distillation training method, etc.

[0115] Furthermore, the ViT distillation training device in this embodiment may also include:

[0116] The input interface 103 is used to obtain the ViT distillation training program imported from the outside, and save the obtained ViT distillation training program to the memory 102. It can also be used to obtain various instructions and parameters transmitted by the external terminal device, and transmit them to the processor 101, so that the processor 101 can use the above various instructions and parameters to carry out corresponding processing. In this embodiment, the input interface 103 can specifically include but is not limited to a USB interface, a serial interface, a voice input interface, a fingerprint input interface, a hard disk reading interface, etc.

[0117] The output interface 104 is used to output various data generated by the processor 101 to a terminal device connected thereto, so that other terminal devices connected to the output interface can obtain various data generated by the processor 101. In this embodiment, the output interface 104 may specifically include but is not limited to a USB interface, a serial interface, and the like.

[0118] The communication unit 105 is used to establish a remote communication connection between the server operation service optimization configuration device and the external server, so that the ViT distillation training device can mount the image file to the external server. In this embodiment, the communication unit 105 may specifically include but is not limited to a remote communication unit based on wireless communication technology or wired communication technology.

[0119] The keyboard 106 is used to obtain various parameter data or instructions input by the user by tapping the keycaps in real time.

[0120] The display 107 is used to display relevant information of the ViT distillation training process in real time.

[0121] The mouse 108 can be used to assist the user in inputting data and simplify the user's operation.

[0122] The present invention also discloses a readable storage medium, wherein the readable storage medium includes a random access memory (RAM), a memory, a read-only memory (ROM), an electrically programmable ROM, an electrically erasable programmable ROM, a register, a hard disk, a removable hard disk, a CD-ROM, or any other form of storage medium known in the technical field. The readable storage medium stores a ViT distillation training program, and when the ViT distillation training program is executed by a processor, the steps of the ViT distillation training method described in any one of the above items are implemented.

[0123] In summary, the present invention effectively improves the classification performance and generalization ability of the ViT model on long-tail datasets by combining long-tail dataset preprocessing, teacher network training, knowledge distillation, and loss reweighting techniques.

[0124] In this specification, each embodiment is described in a progressive manner, and each embodiment focuses on the differences from other embodiments. The same or similar parts between the embodiments can be referred to each other. For the method disclosed in the embodiment, since it corresponds to the system disclosed in the embodiment, the description is relatively simple, and the relevant parts can be referred to the method part.

[0125] Professionals may further appreciate that the units and algorithm steps of each example described in conjunction with the embodiments disclosed herein can be implemented in electronic hardware, computer software, or a combination of the two. In order to clearly illustrate the interchangeability of hardware and software, the composition and steps of each example have been generally described in the above description according to function. Whether these functions are performed in hardware or software depends on the specific application and design constraints of the technical solution. Professionals and technicians may use different methods to implement the described functions for each specific application, but such implementation should not be considered to be beyond the scope of the present invention.

[0126] In the several embodiments provided by the present invention, it should be understood that the disclosed systems, systems and methods can be implemented in other ways. For example, the system embodiments described above are only schematic. For example, the division of the units is only a logical function division. There may be other division methods in actual implementation, such as multiple units or components can be combined or integrated into another system, or some features can be ignored or not executed. Another point is that the mutual coupling or direct coupling or communication connection shown or discussed can be through some interfaces, indirect coupling or communication connection of systems or units, which can be electrical, mechanical or other forms.

[0127] The units described as separate components may or may not be physically separated, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed on multiple network units. Some or all of the units may be selected according to actual needs to achieve the purpose of the solution of this embodiment.

[0128] In addition, each functional module in each embodiment of the present invention may be integrated into one processing unit, or each module may exist physically separately, or two or more modules may be integrated into one unit.

[0129] Similarly, each processing unit in each embodiment of the present invention may be integrated into one functional module, or each processing unit may exist physically, or two or more processing units may be integrated into one functional module.

[0130] The steps of the method or algorithm described in conjunction with the embodiments disclosed herein may be implemented directly using hardware, a software module executed by a processor, or a combination of the two. The software module may be placed in a random access memory (RAM), a memory, a read-only memory (ROM), an electrically programmable ROM, an electrically erasable programmable ROM, a register, a hard disk, a removable disk, a CD-ROM, or any other form of storage medium known in the art.

[0131] Finally, it should be noted that, in this article, relational terms such as first and second, etc. are only used to distinguish one entity or operation from another entity or operation, and do not necessarily require or imply any such actual relationship or order between these entities or operations. Moreover, the terms "include", "comprise" or any other variants thereof are intended to cover non-exclusive inclusion, so that a process, method, article or device including a series of elements includes not only those elements, but also other elements not explicitly listed, or also includes elements inherent to such process, method, article or device. In the absence of further restrictions, the elements defined by the sentence "comprise a ..." do not exclude the presence of other identical elements in the process, method, article or device including the elements.

[0132] The above is a detailed introduction to the ViT distillation training method, system, device and readable storage medium provided by the present invention. Specific examples are used in this article to illustrate the principles and implementation methods of the present invention. The description of the above embodiments is only used to help understand the method of the present invention and its core idea. It should be pointed out that for ordinary technicians in this technical field, without departing from the principles of the present invention, several improvements and modifications can be made to the present invention, and these improvements and modifications also fall within the scope of protection of the claims of the present invention.

Claims

1. A ViT distillation training method, characterized in that: include: Select a long-tail data set and preprocess the long-tail data set; Train the convolutional neural network to generate a teacher model; Generate data samples based on the long-tail dataset, input them into the teacher model, and generate soft labels; Build a student model based on the visual Transformer and configure the knowledge distillation function and loss function for the student model; Generate data samples based on the long-tail data set, input them into the student model, generate classification results, and optimize the student model parameters using the loss function based on the classification results and soft labels to train the student model; After the training is completed, the student model is verified and output.

2. The ViT distillation training method according to claim 1, characterized in that: The selecting of the long-tail data set and preprocessing of the long-tail data set include: Choose CIFAR-10-LT or ImageNet-LT as the long-tail dataset; Oversample the tail categories of the long-tail dataset and undersample the head categories of the long-tail dataset to generate data samples; Adjust the class weights according to the number of data samples in each class and reweight the loss of the tail classes when calculating the loss function; All sample data are weakly enhanced, and sample data of tail categories are strongly enhanced to generate out-of-distribution images.

3. The ViT distillation training method according to claim 2, characterized in that: The convolutional neural network is trained to generate a teacher model, including: Train the convolutional neural network CSPDarkNet; During training, the model parameters are updated by calculating the SAM perturbation direction and the gradient after the perturbation; After training is completed, a teacher model is generated.

4. The ViT distillation training method according to claim 3, characterized in that: The method of generating data samples based on the long-tail data set, inputting the data samples into the teacher model, and generating soft labels includes: The data samples after strong enhancement are input into the teacher model to generate soft labels.

5. The ViT distillation training method according to claim 4, characterized in that: The student model is constructed based on the visual Transformer, and the knowledge distillation function and loss function are configured for the student model, including: Use visual Transformer as the student model; DS Token is introduced into the student model for distillation learning, and DS Token is used to learn features by imitating the soft labels of the teacher model; Configure the loss function for calculating classification loss and distillation loss for the student model; The classification loss is calculated using cross entropy loss, and the distillation loss is calculated by minimizing the KL divergence function.

6. The ViT distillation training method according to claim 5, characterized in that: The method of optimizing the student model parameters by using a loss function based on the classification results and the soft labels includes: Based on the classification results and soft labels, the classification loss and distillation loss are calculated using the loss function to jointly optimize the parameters of the student model; When calculating the distillation loss, the weight of the soft labels of the tail category sample data is increased to ensure the attention of the tail category sample data.

7. The ViT distillation training method according to claim 2, characterized in that: The weak enhancement processing includes: random cropping processing, horizontal flipping processing and color adjustment processing; the strong enhancement processing includes: random rotation processing, cropping processing with a preset amplitude and color dithering processing.

8. A ViT distillation training system, characterized in that: The system adopts the ViT distillation training method according to any one of claims 1 to 7; The system comprises: The preprocessing module is used to select and preprocess the long-tail data set; the teacher model building module is used to train the convolutional neural network and generate the teacher model; The soft label generation module is used to generate data samples based on the long-tail data set, input the teacher model, and generate soft labels; The student model building module is used to build a student model based on the visual Transformer and configure the knowledge distillation function and loss function for the student model; The model training module is used to generate data samples based on the long-tail data set, input the data into the student model, generate classification results, and optimize the student model parameters using the loss function based on the classification results and soft labels to train the student model; The model verification output module is used to verify and output the student model after training is completed.

9. A ViT distillation training device, characterized in that: include: A memory for storing the ViT distillation training program; A processor, configured to implement the steps of the ViT distillation training method as described in any one of claims 1 to 7 when executing the ViT distillation training program.

10. A readable storage medium, characterized in that: The readable storage medium stores a ViT distillation training program, and when the ViT distillation training program is executed by the processor, the steps of the ViT distillation training method according to any one of claims 1 to 7 are implemented.

Images