Image classification method and system based on pruning deep learning model

Abstract

The application discloses an image classification method and system based on a pruned deep learning model, wherein the method comprises the following steps: acquiring a deep learning model, performing pruning operation on the deep learning model to obtain a pruned deep learning model; the pruning operation comprises the following steps: adopting a genetic algorithm to perform pruning operation on a multi-head attention mechanism layer Multi-head Attention and a fully connected feedforward network FFN of each encoder Transformer Encoder in a Vision Transformer model; obtaining a trained pruned deep learning model; acquiring an image to be classified, inputting the image to be classified into the trained pruned deep learning model, and obtaining a classification result of the image.

Description

Technical Field

[0001] This invention relates to the field of deep learning model pruning and compression technology, and in particular to an image classification method and system based on pruning deep learning models. Background Technology

[0002] In the past, deep learning models have achieved remarkable success in fields such as image recognition and natural language processing. However, these models typically have high computational costs and a large number of parameters, limiting their operation on mobile devices and hindering their application and development across various industries. To address these issues, researchers have proposed various model compression techniques aimed at reducing the number of parameters and computational requirements of models while maintaining performance as much as possible. Pruning is one such method, which reduces model complexity by removing unimportant weights, thereby reducing the number of parameters and computational load. This method can effectively reduce model size while maintaining accuracy. Quantization is another commonly used compression method, which reduces model storage requirements by decreasing the precision of model weights. This method can reduce model weights from 32-bit floating-point numbers to 8-bit integers, significantly reducing storage space and computational requirements, while in many cases, the impact on model performance is negligible. Although these compression techniques have achieved some success, these results have mainly focused on convolutional neural networks (CNNs).

[0003] With continuous improvements in model architecture, the Transformer model has achieved revolutionary progress in the field of natural language processing. As this model architecture continues to evolve, recent research has begun to explore the application of the Vision Transformer (ViT) based on the Transformer in computer vision tasks, especially large-scale image classification. Compared to traditional convolutional neural networks (CNNs), ViT exhibits significant performance advantages in handling large-scale image classification tasks, as its entire structure does not contain any convolutional operations. DeepViT further optimizes the ViT structure by introducing a knowledge distillation strategy, achieving performance exceeding that of the original ViT using only the ImageNet-1k dataset. Furthermore, DeepViT proposes a simple yet effective re-attention mechanism, further enhancing model performance by reconstructing the attention map. With the continuous development of ViT and its derivative models, their application scope has expanded to multiple image-related task domains, including image segmentation and object detection, demonstrating the broad applicability and potential of the ViT architecture in computer vision.

[0004] While Vision Transformer (ViT) performs exceptionally well on large-scale classification tasks, it still suffers from high complexity, primarily stemming from its processing flow and attention mechanism. First, ViT segments the input image into patches; then, these patches are converted into tokens, each corresponding to an element in the input sequence. Since ViT uses self-attention to process these tokens, it needs to calculate the relationship between each token and all other tokens in the sequence. This process involves a large number of parameters and computations, and the computational complexity increases quadratically, especially as the number of tokens increases. Furthermore, the depth of the ViT model also adds to the computational burden, as each layer needs to process and update the feature representations of all tokens. Therefore, while achieving powerful image classification capabilities, ViT also faces the challenges of large model size and high computational cost. Further compression of the model to reduce its runtime resource consumption and enable wider application has become a practical need. Summary of the Invention

[0005] To address the issues of excessive parameters and computational resource consumption in the Vision Transformer model, this invention proposes an image classification method and system based on a pruned deep learning model, denoted as MSCViT. This invention reduces the computational cost and number of parameters by pruning the token, attention head, and FFN modules of the model. The importance evaluation method for the token module is improved to reduce the need for large amounts of additional computational resources, and a hierarchical pruning method based on intelligent optimization algorithms is employed for the attention head and FFN modules. Through these methods, the number of parameters and FLOPs of the model are effectively reduced, while maintaining the accuracy or only slightly decreasing it.

[0006] On the one hand, image classification methods based on pruned deep learning models are provided, including:

[0007] A deep learning model is obtained, and a pruning operation is performed on the deep learning model to obtain a pruned deep learning model; wherein, the deep learning model is a Vision Transformer model; the pruning operation includes: performing a genetic algorithm to prune the Multi-head Attention layer and the fully connected feedforward network FFN of each encoder Transformer Encoder in the Vision Transformer model; the pruning operation also includes: adding a Token Select layer between the Multi-head Attention layer and the fully connected feedforward network FFN of each encoder Transformer Encoder, wherein the Token Select layer calculates an importance score for each token and performs the pruning operation based on the importance score to reduce the number of tokens;

[0008] Construct a training set, which consists of images with known image classification labels; input the training set into a deep learning model, train the model, and obtain a trained deep learning model; copy the model parameters of the trained deep learning model into a pruned deep learning model to obtain a trained pruned deep learning model.

[0009] The image to be classified is obtained and then input into the trained and pruned deep learning model to obtain the classification result.

[0010] On the other hand, an image classification system based on a pruned deep learning model is provided, including:

[0011] The pruning module is configured to: acquire a deep learning model, perform pruning operations on the deep learning model, and obtain a pruned deep learning model; wherein the deep learning model is a Vision Transformer model; the pruning operation includes: performing pruning operations on the Multi-head Attention layer and the fully connected feedforward network FFN of each encoder Transformer Encoder in the Vision Transformer model using a genetic algorithm; the pruning operation further includes: adding a Token Select layer between the Multi-head Attention layer and the fully connected feedforward network FFN of each encoder Transformer Encoder, wherein the Token Select layer calculates an importance score for each token and performs pruning operations based on the importance score to reduce the number of tokens;

[0012] The training module is configured to: construct a training set, wherein the training set consists of images with known image classification labels; input the training set into a deep learning model, train the model, and obtain a trained deep learning model; and copy the model parameters of the trained deep learning model into a pruned deep learning model to obtain a trained pruned deep learning model.

[0013] The classification module is configured to: acquire the image to be classified, input the image to be classified into the trained and pruned deep learning model, and obtain the classification result of the image.

[0014] Furthermore, an electronic device is also provided, including:

[0015] Memory, used for non-transitory storage of computer-readable instructions; and

[0016] Processor, for executing the computer-readable instructions,

[0017] When the computer-readable instructions are executed by the processor, they perform the method described in the first aspect above.

[0018] In another aspect, a storage medium is also provided for non-transitory storage of computer-readable instructions, wherein when the non-transitory computer-readable instructions are executed by a computer, the method described in the first aspect is performed.

[0019] In another aspect, a computer program product is also provided, including a computer program that, when run on one or more processors, is used to implement the method described in the first aspect above.

[0020] The above technical solution has the following advantages or beneficial effects:

[0021] (1) Token pruning simplifies the core structure of the Vision Transformer (ViT) model by reducing the number of tokens involved in self-attention computation, thereby reducing the model's parameter count and computational complexity. Since ViT segments the input image into multiple tokens, each token representing a small patch of the image, pruning removes tokens that contribute little to the final classification or recognition task, reducing unnecessary information processing and thus alleviating the model's computational burden. This method not only reduces the model's storage and inference costs but also improves the model's generalization ability by removing noise and unimportant features, avoiding overfitting. Furthermore, the pruned model, with fewer parameters, may be more stable during training, helping it learn more robust feature representations. In addition, pruning enhances the model's interpretability because retaining fewer but more important tokens allows researchers to more intuitively understand the model's decision-making process, enabling them to deploy and use these optimized models with greater confidence in practical applications.

[0022] (2) The proposed method of using a genetic algorithm to search for the optimal pruning configuration of different layers in the Attention head and FFN module demonstrates significant beneficial effects. By employing targeted pruning rates to prune different layers of the neural network, the sparsity of the model can be precisely adjusted, optimizing performance and efficiency. This method fully utilizes the different importance of each layer in feature extraction, allowing more parameters to be retained for key layers to capture important information, while less important layers are pruned more aggressively to reduce redundancy. This differentiated approach not only improves the model's generalization ability and avoids overfitting, but also balances the use of computational resources, enabling the model to maintain accuracy while reducing inference time and memory usage, thus enhancing the model's practicality and flexibility in resource-constrained environments.

[0023] (3) The weight sharing mechanism in this invention greatly reduces the time required to retrain each compressed model and significantly improves the search efficiency of the genetic algorithm. Attached Figure Description

[0024] The accompanying drawings, which form part of this invention, are used to provide a further understanding of the invention. The illustrative embodiments of the invention and their descriptions are used to explain the invention and do not constitute an improper limitation of the invention.

[0025] Figure 1 Here is a diagram of the Vision Transformer model structure;

[0026] Figure 2 This is a schematic diagram of the attention weighting analysis method of the present invention;

[0027] Figure 3This is a schematic diagram of the token selection layer of the present invention;

[0028] Figure 4 This is a schematic diagram of token fusion pruning in this invention;

[0029] Figure 5 This is a schematic diagram of the genetic algorithm search and pruning configuration of the present invention;

[0030] Figure 6 This is a schematic diagram of the internal connection relationship of the fully connected feedforward network (FFN) of the present invention. Detailed Implementation

[0031] It should be noted that the following detailed descriptions are exemplary and intended to provide further illustration of the invention. Unless otherwise specified, all technical and scientific terms used herein have the same meaning as commonly understood by one of ordinary skill in the art to which this invention pertains.

[0032] Example 1

[0033] This embodiment provides an image classification method based on a pruned deep learning model;

[0034] Image classification methods based on pruned deep learning models include:

[0035] S101: Obtain a deep learning model and perform pruning operations on the deep learning model to obtain a pruned deep learning model; wherein, the deep learning model is a Vision Transformer model; the pruning operation includes: performing pruning operations on the multi-head attention mechanism layer and the fully connected feedforward network FFN of each encoder Transformer Encoder in the Vision Transformer model using a genetic algorithm; the pruning operation further includes: adding a Token Select layer between the multi-head attention mechanism layer and the fully connected feedforward network FFN of each encoder Transformer Encoder, wherein the Token Select layer calculates an importance score for each token and performs pruning operations based on the importance score to reduce the number of tokens;

[0036] S102: Construct a training set, which consists of images with known image classification labels; input the training set into a deep learning model, train the model, and obtain a trained deep learning model; copy the model parameters of the trained deep learning model into a pruned deep learning model to obtain a trained pruned deep learning model.

[0037] S103: Obtain the image to be classified, input the image to be classified into the trained and pruned deep learning model, and obtain the classification result of the image.

[0038] Furthermore, such as Figure 1 As shown, the Vision Transformer model includes:

[0039] The layers are connected in sequence: Linear Projection of Flattened Patches, Position Encoding Layer, L Transformer Encoders in series, Multilayer Perceptron, and Classification Layer.

[0040] Furthermore, the Vision Transformer model includes:

[0041] By dividing the input image into blocks;

[0042] The flattening layer stretches n×n image patches into a one-dimensional vector;

[0043] The positional encoding layer performs a linear projection on a one-dimensional vector and adds positional encoding to obtain several tokens;

[0044] L cascaded Transformer Encoders perform L feature extractions on the token to obtain the extracted features;

[0045] The extracted features are input into a multilayer perceptron for computation, and the classification layer provides classification labels based on the computation results of the multilayer perceptron.

[0046] For example, a flattening layer works by segmenting the input image into fixed-size image patches, flattening each image patch into a one-dimensional vector, and then mapping the vector to a high-dimensional embedding space through a linear projection (convolutional layer).

[0047] For example, the position coding layer generates position codes PE using sine and cosine functions, as shown in the following formula:

[0048]

[0049]

[0050] Where pos is the position index, i is the dimension index, and d is the position index. model It is the embedding dimension, and 10000 is a constant used to control the frequency of the sine and cosine functions.

[0051] The Vision Transformer model consists of L encoders connected in series. Each encoder includes:

[0052] The first layer normalization module LayerNorm, the multi-head attention mechanism layer Multi-headAttention, the first adder, the second layer normalization module LayerNorm, the fully connected feedforward network FFN, and the second adder are connected in sequence.

[0053] The input of the first layer normalization module LayerNorm is the input of the current encoder Transformer Encoder; the output of the second adder is the output of the current encoder Transformer Encoder.

[0054] The input of the first-level normalization module LayerNorm is connected to the input of the first adder;

[0055] The output of the first adder is connected to the input of the second adder.

[0056] For example, the first layer normalization module LayerNorm is used to perform normalization processing on the input data, so that the features of each sample have the same mean and standard deviation within the same layer.

[0057] Furthermore, the pruned deep learning model includes: a flattening layer, a position encoding layer, L cascaded encoders (Transformer Encoders), a multilayer perceptron, and a classification layer, connected in sequence.

[0058] The difference between the pruned deep learning model and the Vision Transformer model lies in:

[0059] In the pruned deep learning model, a Token Select layer is added between the Multi-head Attention layer and the fully connected feedforward network FFN in each encoder (Transformer Encoder). The input of the Token Select layer is connected to the output of the Multi-head Attention layer, and the output of the Token Select layer is connected to the input of the fully connected feedforward network FFN.

[0060] There is no Token Select layer between the Multi-head Attention layer of each encoder in the Vision Transformer model and the fully connected feedforward network FFN.

[0061] The difference between the pruned deep learning model and the Vision Transformer model also lies in:

[0062] In the pruning deep learning model, the multi-head attention mechanism layer of each encoder (Transformer Encoder) and the fully connected feedforward network (FFN) are pruned using a genetic algorithm.

[0063] Furthermore, the Token Select layer calculates an importance score for each token and performs pruning operations based on the importance score to reduce the number of tokens, including:

[0064] (1-1) Calculate the importance score for each token;

[0065] (1-2) Compare the importance score of each token with the set score threshold, merge tokens that are below the score threshold, and retain tokens that are above the score threshold.

[0066] Further, step (1-1) calculates the importance score for each token, specifically including:

[0067] In Vision Transformer (ViT), attention weights reflect the strength of the interaction between different tokens;

[0068] In Vision Transformer, the input sequence X is mapped to three matrices: Query, Key, and Value. In each Transformer Encoder, the attention weight A of each head in the multi-head attention mechanism is calculated as shown in formula (2):

[0069]

[0070] Among them, Q h , V h The query matrix, key matrix, and value matrix are respectively converted from the token input in each header h. k This is the scaling factor;

[0071] The attention allocation matrix for a single token to other tokens is derived as follows:

[0072]

[0073] Among them, A h [i,j] represents the attention weight of the i-th token to the j-th token. This represents the query value of the i-th token in the h-th attention head. This represents the key-value pair of the j-th token in the h-th attention header. This represents the key value of the m-th token in the h-th attention head.

[0074] For the i-th token in the h-th attention head, the attention weights between the i-th token and all other tokens are calculated: using the query vector of the i-th token. key vector of each token Performing a dot product, followed by scaling and softmax normalization, yields a probability distribution, which is the attention allocation matrix A. h [i,j]; The probability distribution represents the attention allocation of the i-th token to all other tokens.

[0075] The attention allocation matrix obtained by formula (3) reflects the global importance. The importance score of the token within the head is obtained by the attention allocation matrix. Within a single attention head, the importance score of the i-th token is calculated by averaging the attention weight of the i-th token to all tokens (including itself).

[0076]

[0077] in, This represents the initial importance score of the i-th token in the h-th attention head. The initial importance score of the i-th token reflects the overall importance of the i-th token in the current attention head. The larger the value, the stronger the interaction between the i-th token and other tokens, and the more important its role in feature extraction.

[0078] Since the model is based on a multi-head self-attention mechanism, it means that there are multiple attention heads within each Transformer Encoder. The initial importance scores of the same token in all attention heads within the Transformer Encoder are aggregated and averaged to obtain the final importance score of the i-th token in the entire Transformer Encoder.

[0079]

[0080] in, This represents the final importance score of the i-th token in the current Transformer Encoder. This aggregation takes into account the possibility that different attention heads may focus on different types of features, comprehensively reflecting the importance of the token across multiple feature subspaces.

[0081] The importance score for each token is represented as:

[0082]

[0083] Among them, Imp i This represents the importance score of the i-th token.

[0084] like Figure 2 As shown, this invention provides a method for evaluating the importance of tokens based on attention weight analysis. The relevant computational information can be directly obtained from the attention graph calculated in the self-attention module. This method avoids introducing additional parameters and is computationally efficient.

[0085] Furthermore, steps (1-2) involve fusing tokens that are below a score threshold. The specific fusing process includes:

[0086] (1-2-1) Identify several tokens whose importance scores are below a set threshold based on the importance scores of the tokens;

[0087] (1-2-2) Perform a weighted summation operation on several tokens whose importance scores are below the threshold to generate a merged token;

[0088] The specific formula is as follows:

[0089] x fused =∑ i∈N a i x i

[0090] Where, x fusedIt is the merged token, N is a set of several tokens with importance scores below a threshold, and a i It is the weight corresponding to the token at index i, x i It is the token at index i.

[0091] (1-2-3) Add the calculated merged token to the set of tokens whose importance score is higher than the set threshold.

[0092] It should be understood that, after the token importance calculation is completed in the selected Transformer Encoder, in order to implement token pruning, a Token Selection layer (T-Select) is added between the MHSA and FFN modules in each Transformer layer, such as... Figure 3 As shown, this operation does not add any extra parameters. The Token selection layer receives the input feature Z. l Then, based on the given pruning ratio v1, select the most important (1-v) branches to retain. l N tokens are selected, and the least important v1N tokens are merged based on importance scores calculated using an attention weighting analysis method. After merging tokens below a certain score threshold, the token length of subsequent layers is reduced, and subsequent token selection layers choose from the tokens retained by the previous layer. Figure 4 As shown.

[0093] Furthermore, in the Vision Transformer model, the multi-head attention mechanism layer and the fully connected feedforward network FFN of each encoder are pruned using a genetic algorithm, including:

[0094] (2-1): Genetic algorithm is used to search for the optimal pruning strategy for the Multi-head Attention layer and the fully connected feedforward network FFN of each Transformer Encoder;

[0095] (2-2) The optimal pruning strategy found in the search is used to perform pruning operations.

[0096] Furthermore, such as Figure 5 As shown, in (2-1): a genetic algorithm is used to search for the optimal pruning strategy for the Multi-head Attention layer and the fully connected feedforward network FFN of each Transformer Encoder, including:

[0097] (2-1-1) For different pruning rates, configure p i =[r1,r2,...,r n ,s1,s2,...,s n ], where r1, r2, ..., r n s1, s2, ..., s represents the pruning ratio for the first to nth multi-head attention mechanism layers. n This represents the pruning ratio for the 1st to nth fully connected feedforward networks (FFNs). The search objective is as follows:

[0098] Obj(p i )=αAccuracy(p i )+βΔFLOPs(p i (7)

[0099]

[0100] Where TP is the number of samples correctly predicted as positive, TN is the number of samples correctly predicted as negative, FP is the number of negative samples incorrectly predicted as positive, and FN is the number of positive samples incorrectly predicted as negative.

[0101]

[0102] in, This indicates the number of floating-point operations performed on the original model. This indicates that the pruning rate is configured as p i The number of floating-point operations in the pruned model.

[0103] Accuracy(p i ) indicates that in the pruning strategy p i The accuracy of the model on the test set, ΔFLOPs(p i ) indicates that in the pruning strategy p i The difference between the model's FLOPs and the target FLOPs, obtained by a dedicated calculation module, is used to calculate the objective function Obj(p). i )maximize;

[0104] The dedicated calculation module uses the following formula:

[0105]

[0106] Where, d model It is a fixed dimension of the projection of each patch through a linear layer, N is the number of patches obtained from the segmentation, and d patch L is the dimension of each patch, and L is the number of Transformer Encoders.

[0107] α and β represent the proportions of accuracy and FLOPs difference when evaluating different pruning rate configurations, respectively. The larger α is, the higher the proportion of accuracy factor, while the larger β is, the higher the proportion of FLOPs difference factor. α + β = 1.

[0108] (2-1-2) Define the search space as follows:

[0109] O∈(0.001,0.5)(8)

[0110] O indicates that the pruning ratio is to find the optimal value between 0.001 and 0.5. Through actual analysis, if the pruning ratio is too high, it will cause a significant decrease in accuracy. Therefore, the pruning ratio is limited to between 0.001 and 0.5, which is a balance between accuracy and model computation.

[0111] (2-1-3) Within the search space, a set of candidate pruning strategies is randomly initialized. Each candidate pruning strategy is treated as an individual, representing a set of candidate solutions. All candidate solutions are combined into a population set. Each population set is iteratively updated with new strategies. The candidate pruning strategies are represented as follows:

[0112] Pol FFN ={r1,r2,...,r 11 ,r 12} (9)

[0113] Pol head ={s1,s2,...,s 11 ,s 12} (10)

[0114] Pol represents the result of the candidate pruning strategy, and the right side of equation (9) is {r1,r2,...,r 11 ,r 12} represents the pruning ratio of the fully connected feedforward network (FFN) in the Transformer Encoder from layers 1 to 12, denoted as Pol. FFN The right side of equation (10) [s1,s2,...,s 11 ,s 12 ] represents the pruning ratio of the multi-head attention mechanism layers in the Transformer Encoder from layers 1 to 12, denoted as Pol. head ;

[0115] (2-1-4) The fitness function is set in the same way as the objective function, and the expression is as follows:

[0116] Obj(p i)=αAccuracy(p i )+βΔFLOPs(p i );

[0117] (2-1-5) After defining the fitness function, the initial candidate pruning strategy evolves to generate a new generation of population through selection, crossover, and mutation operations; the evolution process adopts a roulette wheel selection mechanism, completely replacing the original population with offspring. After 50 iterations, the algorithm converges to the optimal pruning strategy.

[0118] Furthermore, step (2-2) employs the optimal pruning strategy found during the search for pruning operations, including:

[0119] Pruning methods for fully connected feedforward networks (FFNs): such as Figure 6 As shown, the fully connected feedforward network FFN is known to include an input layer, an intermediate layer, and an output layer connected sequentially. In the intermediate layer of the fully connected feedforward network FFN, the number of neurons is reduced through pruning. For the l-th layer of the Transformer Encoder in the fully connected feedforward network FFN, feature Z is... l The features of each neuron j are: The intermediate layer is implemented through a hidden layer;

[0120] By calculating the Gram matrix Dependency score Assess the importance score of each neuron;

[0121] Gram matrix The elements in are defined as Here, k is a kernel function used to calculate the similarity between two vectors.

[0122] Dependency score The formula is:

[0123]

[0124] in, It is defined in The Gram matrix on the matrix, tr represents the trace of the matrix, and CLC is a constant matrix; the dependency score is the importance score of each neuron.

[0125] The neuron pruning ratio k obtained from the genetic algorithm search l Retain the one with the highest importance score One neuron was removed; the rest were deleted.

[0126] Furthermore, the pruning method for the Multi-head Attention mechanism layer:

[0127] In the Multi-head Attention layer of the l-th Transformer Encoder, the output feature of each attention head h is represented as...

[0128] Calculate the Gram matrix K for each attention head. l Dependency score Assess the importance score for each head;

[0129] Based on the pruning ratio σ obtained from the search l The top σ scorers will be selected. l ·N h Keep one focus item and delete the rest. N h The total number of attention heads.

[0130] It should be understood that using a genetic algorithm to search for different pruning ratios for each layer of the Multi-head Attention layer, for each individual attention head and FFN, allows for targeted pruning. This targeted pruning helps find a better balance between reducing computational burden and maintaining model performance, allowing the model to maintain high performance in critical layers while reducing computational resource usage in layers with less impact on the final output.

[0131] In Vision Transformer, the single attention head of the Multi-head Attention layer and the FFN module are also important components and play a role. For these two modules, a genetic algorithm is used to find the pruning strategy of each layer in the FFN and the single attention head of the Multi-head Attention layer. Under the constraint of the pruning range, the ViT network after compression by the single attention head of the FFN and the Multi-head Attention layer obtains the optimal solution by balancing the computational cost, accuracy and inference time.

[0132] By copying the model parameters of the trained deep learning model into the pruned deep learning model, this invention uses an improved shared-weight model accuracy evaluation method for fast accuracy evaluation during search, without having to train each pruned model from scratch.

[0133] During the search process, instead of retraining the model for each pruning strategy, we efficiently evaluate model accuracy through weight sharing. We pre-train the complete model weights W, allowing us to predict the accuracy of independently trained sub-models based on the weights inherited by each sub-model. Thanks to this sharing mechanism, we only need to train the model once to obtain the corresponding weights, which can then be used to evaluate the accuracy of various pruning strategies. This approach significantly reduces the time required compared to retraining each compressed model and substantially improves the search efficiency of the genetic algorithm.

[0134] This embodiment provides a compressed retraining method based on a modular compression framework, which is divided into two stages: 1) Pruning stage. This involves removing unimportant weights or neurons from the ViT model. For the token module, the importance of the token is evaluated based on the attention weight analysis method, and fusion pruning is performed. For the Attention head and FFN modules, the importance of each neuron is evaluated by calculating the Gram matrix and dependency score, and hierarchical pruning is performed based on the searched pruning rate. The purpose of this stage is to reduce the complexity and computational cost of the model while preserving as much of the model's key information as possible. 2) Retraining stage is performed after pruning. The ILSVRC-2012 subset of the ImageNet dataset is used, and the AdamW optimizer and cosine decay strategy are employed to fine-tune the remaining weights to recover some of the performance loss that may have been caused by pruning, ensuring that the model maintains high accuracy and generalization ability while reducing parameters.

[0135] First, the importance score of a token is calculated by averaging the attention weights of each token to all tokens (including itself) within the selected Transformer Encoder. After the token importance calculation is completed in the selected Transformer Encoder, a Token Selection layer (T-Select) is added between the MHSA and FFN modules in each Transformer layer to implement token pruning. This operation does not add any additional parameters. T-Select receives the input feature Z. l Then, according to the given pruning ratio v l Select and retain the most important (1-v) l N tokens are used, and the least important token is merged based on its importance score. l There are N tokens, and their scores are calculated using an attention weighting analysis method. Finally, tokens below a score threshold are merged, reducing the token length of subsequent layers, and subsequent T-Select operations select from the tokens retained in the previous layer.

[0136] For the single attention head and FFN module of the Multi-head Attention mechanism layer, a genetic algorithm is first used to search for the optimal pruning strategy for different layers. A set of pruning strategies is randomly initialized in the search space, and the fitness function and objective function are set. The initial candidate pruning strategies evolve to generate a new generation of population through selection, crossover, and mutation operations. This process uses a roulette wheel selection mechanism to completely replace the original population with offspring. After 50 iterations, the algorithm converges to an approximately optimal pruning strategy. In the intermediate layer of the FFN module, the pruning ratio k of the neurons obtained from the search is... l Keep the most important 1 neuron. By calculating the Gram matrix Dependency score To evaluate the importance of each neuron, the neuron with the highest score is retained, and the rest are pruned. For a single attention head in the Multi-head Attention layer, in the Multi-head Self-Attention (MHSA) of the l-th layer, the output feature of each head h is represented as Z. l Based on the pruning ratio σ obtained from the search... l The output features of each head are simplified by average pooling to form... Then calculate the Gram matrix K for each head. l Dependency score This is used to assess the importance of each neuron. The neuron with the highest score is retained, while the rest are pruned.

[0137] The pruned model was fine-tuned to restore accuracy. Training was performed using the PyTorch framework on a specific GPU, following the initial configuration of the ViT model. The AdamW optimizer and cosine decay strategy were employed to achieve more stable learning rate adjustments during training. Simultaneously, the learning rate and batch size were reduced to allow for finer control over the model's convergence process, ensuring better generalization ability during training. After initial convergence, the pruned ViT model underwent 100 and 80 iterations of fine-tuning respectively to further improve its performance.

[0138] The tokens module uses an attention weight analysis method to evaluate the importance of each token and prunes and removes unimportant tokens based on their importance ranking, thereby reducing the burden on the model in this part.

[0139] The Multi-head Attention mechanism employs an intelligent optimization algorithm to search for the optimal pruning strategy for each layer in the model, thereby achieving precise pruning at different layers.

[0140] The pruning method of the FFN module uses a genetic algorithm to search for the optimal pruning strategy for each layer in the model, achieving different pruning rates for different layers.

[0141] By fine-tuning and retraining, the model's performance can be restored to the level before pruning, as much as possible.

[0142] Attention weight analysis leverages the inherent characteristics of the attention mechanism to calculate the attention magnitude between a single token and other tokens as a criterion for importance assessment.

[0143] The removal of unimportant tokens is not done by direct deletion, but by merging them into a single token based on their characteristics, thereby reducing the number of tokens.

[0144] Based on the characteristics of FFN module neurons, a genetic algorithm is used to search for the optimal pruning ratio of each FFN module in the ViT model to adapt to different levels of the model.

[0145] Based on the characteristics of a single attention head in the Multi-head Attention layer, a genetic algorithm is used to search for the optimal pruning ratio of a single attention head in each Multi-head Attention layer of the ViT model, so as to adapt to the distribution of the number of single attention heads in the Multi-head Attention layer in different layers of the model.

[0146] The accuracy of the model is efficiently evaluated through weight sharing. The complete model weights W are pre-trained, and the accuracy of each sub-model after independent training can be predicted based on the weights inherited by the sub-models, improving the search efficiency of the genetic algorithm.

[0147] The AdamW optimizer and cosine decay strategy are employed to achieve more stable learning rate adjustments during training. Simultaneously, the learning rate and batch size are reduced to allow for more precise control over the model's convergence process.

[0148] This invention discloses a modular compression method for Vision Transformer (ViT), relating to the field of machine learning. It aims to reduce the computational and parameter requirements of the model through pruning, adapting it to resource-constrained environments. The method is implemented in two steps: pruning the ViT model and fine-tuning the pruned model to restore performance. ViT model pruning involves three modules: tokens, a single attention head in the Multi-head Attention layer, and the FFN (Fast Forwarding Network). For the token module, attention weight analysis is used to evaluate token importance, and token fusion is employed to reduce the number of tokens. For the single attention head and FFN modules in the Multi-head Attention layer, a genetic algorithm is used to search for the optimal pruning configuration for each layer. Weight sharing is also used during the search process to rapidly improve search speed. The importance of each neuron is evaluated based on the Gram matrix and dependency score; neurons below a threshold are pruned. Fine-tuning of the pruned model employs the AdamW optimizer and cosine decay strategy, with a low learning rate and batch size, to better control the model's convergence process and restore the model's performance as close as possible to its pre-pruning level. The technical solution of this invention reduces the computational load and parameter count of the Vision Transformer model, reduces the computational resources required during model training, accelerates the inference speed of the model, makes the ViT model easier to deploy on different platforms, and reduces computational and storage costs.

[0149] Example 2

[0150] This embodiment provides an image classification system based on a pruned deep learning model, including:

[0151] The pruning module is configured to: acquire a deep learning model, perform pruning operations on the deep learning model, and obtain a pruned deep learning model; wherein the deep learning model is a Vision Transformer model; the pruning operation includes: performing pruning operations on the Multi-head Attention layer and the fully connected feedforward network FFN of each encoder (TransformerEncoder) in the Vision Transformer model using a genetic algorithm; the pruning operation further includes: adding a Token Select layer between the Multi-head Attention layer and the fully connected feedforward network FFN of each encoder (TransformerEncoder), wherein the Token Select layer calculates an importance score for each token and performs pruning operations based on the importance score to reduce the number of tokens;

[0152] The training module is configured to: construct a training set, wherein the training set consists of images with known image classification labels; input the training set into a deep learning model, train the model, and obtain a trained deep learning model; and copy the model parameters of the trained deep learning model into a pruned deep learning model to obtain a trained pruned deep learning model.

[0153] The classification module is configured to: acquire the image to be classified, input the image to be classified into the trained and pruned deep learning model, and obtain the classification result of the image.

[0154] It should be noted that the pruning module, training module, and classification module described above correspond to steps S101 to S103 in Embodiment 1. The examples and application scenarios implemented by these modules and their corresponding steps are the same, but they are not limited to the content disclosed in Embodiment 1. It should also be noted that these modules, as part of the system, can be executed in a computer system, such as a set of computer-executable instructions.

[0155] The descriptions of each embodiment in the above embodiments have different focuses. For parts not described in detail in a certain embodiment, please refer to the relevant descriptions in other embodiments.

[0156] The proposed system can be implemented in other ways. For example, the system embodiments described above are merely illustrative, and the division of modules described above is only a logical functional division. In actual implementation, there may be other division methods. For example, multiple modules may be combined or integrated into another system, or some features may be ignored or not executed.

[0157] Example 3

[0158] This embodiment also provides an electronic device, including: one or more processors, one or more memories, and one or more computer programs; wherein, the processor is connected to the memory, and the one or more computer programs are stored in the memory. When the electronic device is running, the processor executes the one or more computer programs stored in the memory to cause the electronic device to perform the method described in Embodiment 1.

[0159] It should be understood that in this embodiment, the processor can be a central processing unit (CPU), or it can be other general-purpose processors, digital signal processors (DSPs), application-specific integrated circuits (ASICs), programmable gate arrays (FPGAs), or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. The general-purpose processor can be a microprocessor or any conventional processor, etc.

[0160] Memory may include read-only memory and random access memory, and provides instructions and data to the processor. A portion of memory may also include non-volatile random access memory. For example, memory may also store information about the device type.

[0161] In the implementation process, each step of the above method can be completed by the integrated logic circuits in the processor hardware or by software instructions.

[0162] The method in Embodiment 1 can be directly implemented by a hardware processor, or implemented by a combination of hardware and software modules within the processor. The software modules can reside in readily available storage media in the art, such as random access memory, flash memory, read-only memory, programmable read-only memory, electrically erasable programmable memory, or registers. This storage medium is located in memory; the processor reads information from the memory and, in conjunction with its hardware, completes the steps of the above method. To avoid repetition, a detailed description is not provided here.

[0163] Those skilled in the art will recognize that the units and algorithm steps described in connection with the various examples of this embodiment can be implemented in electronic hardware or a combination of computer software and electronic hardware. Whether these functions are implemented in hardware or software depends on the specific application and design constraints of the technical solution. Those skilled in the art can use different methods to implement the described functions for each specific application, but such implementation should not be considered beyond the scope of this invention.

[0164] Example 4

[0165] This embodiment also provides a computer-readable storage medium for storing computer instructions, which, when executed by a processor, complete the method described in Embodiment 1.

[0166] The above description is merely a preferred embodiment of the present invention and is not intended to limit the invention. Various modifications and variations can be made to the present invention by those skilled in the art. Any modifications, equivalent substitutions, improvements, etc., made within the spirit and principles of the present invention should be included within the scope of protection of the present invention.

Claims

1. An image classification method based on a pruned deep learning model, characterized in that, The method comprises the following steps: obtaining a deep learning model, performing pruning operation on the deep learning model to obtain a pruned deep learning model; wherein the deep learning model is a Vision Transformer model; the pruning operation comprises: using a genetic algorithm to perform pruning operation on the multi-head attention mechanism layer Multi-head Attention and the fully connected feedforward network FFN of each encoder Transformer Encoder in the Vision Transformer model; the pruning operation further comprises: adding a Token Select layer between the multi-head attention mechanism layer Multi-head Attention and the fully connected feedforward network FFN of each encoder Transformer Encoder, the Token Select layer calculates an importance score for each token, and performs pruning operation according to the importance score to reduce the number of tokens; constructing a training set, the training set being an image with a known image classification label; inputting the training set into the deep learning model to train the model and obtain a trained deep learning model; copying the model parameters of the trained deep learning model into the pruned deep learning model to obtain a trained pruned deep learning model; obtaining an image to be classified, inputting the image to be classified into the trained pruned deep learning model to obtain a classification result of the image. 2.The image classification method based on the pruned deep learning model of claim 1, wherein, The pruned deep learning model comprises: a flattening layer, a position encoding layer, a series of L encoders Transformer Encoder, a multilayer perceptron and a classification layer connected in sequence. The pruned deep learning model and the Vision Transformer model differ in that: a Token Select layer is added between the multi-head attention mechanism layer Multi-head Attention and the fully connected feedforward network FFN of each encoder Transformer Encoder of the pruned deep learning model; the input end of the Token Select layer is connected with the output end of the multi-head attention mechanism layer Multi-head Attention, and the output end of the Token Select layer is connected with the input end of the fully connected feedforward network FFN; The pruned deep learning model and the Vision Transformer model further differ in that: the multi-head attention mechanism layer Multi-head Attention and the fully connected feedforward network FFN of each encoder Transformer Encoder of the pruned deep learning model are both pruned by using a genetic algorithm. 3.The image classification method based on the pruned deep learning model of claim 1, wherein, The Token Select layer calculates an importance score for each token, and performs pruning operation according to the importance score to reduce the number of tokens, comprising: (1-1) calculating an importance score for each token; (1-2) The importance score of each token is compared with a set score threshold, and the tokens below the score threshold are merged, and the tokens above the score threshold are retained. 4.The method of claim 3, wherein the pruning the deep learning model comprises: determining a pruning rate based on the number of the pruning target layers and the total number of layers of the deep learning model; and pruning the deep learning model based on the pruning rate. The importance score of each token is calculated, specifically including: In the Vision Transformer, the input sequence X is mapped to three matrices of queries Query, keys Key and values Value, and in each encoder Transformer Encoder, the attention weight A of each head in the multi-head attention mechanism is calculated as formula (2): wherein Q h , V h are the query matrix, key matrix and value matrix converted from the input tokens token in each head h, respectively, d k is a scaling factor; The attention distribution matrix of a single token to other tokens is calculated as formula (3): wherein A h [i,j] denotes the attention weight of the i-th token token to the j-th token token, denotes the query value of the i-th token token in the h-th attention head, denotes the key value of the j-th token token in the h-th attention head, denotes the key value of the m-th token token in the h-th attention head; In a single attention head, the importance score of the i-th token is calculated by averaging the attention weight of the i-th token to all tokens as formula (4): wherein, represents the initial importance score of the i-th token in the h-th attention head; The initial importance scores of the same token in all attention heads within the encoder Transformer Encoder are aggregated and averaged to obtain the final importance score of the i-th token in the entire encoder Transformer Encoder as formula (5): wherein, denotes the final importance score of the i-th token token in the current encoder Transformer Encoder; the importance score of each token token is denoted as: wherein, Imp i denotes the importance score of the i-th token token. 5.The method of claim 4, wherein the pruning the deep learning model comprises: determining a pruning rate based on the number of the pruning target layers and the total number of layers of the deep learning model; and pruning the deep learning model based on the pruning rate. The tokens below the score threshold are merged, and the specific merging process includes: (1-2-1) According to the importance score of the token, identify several tokens whose importance score is lower than the set threshold; (1-2-2) Perform weighted summation operation on several tokens whose importance score is lower than the threshold to generate a merged token; The specific formula is as follows: x fused =∑ i∈N a i x i wherein x fused is the fused token, N is a set of several tokens token whose importance score is lower than a threshold, a i is the weight of the token corresponding to the index i, x i is the token at index i; (1-2-3) Add the calculated merged token to the token set whose importance score is higher than the set threshold. 6.The method of claim 1, wherein the pruning the deep learning model comprises: determining a pruning rate based on a number of parameters of the deep learning model; and pruning the deep learning model based on the pruning rate. For the multi-head attention mechanism layer Multi-headAttention and the fully connected feedforward network FFN of each encoder Transformer Encoder in the Vision Transformer model, a genetic algorithm is used for pruning operation, including: For the multi-head attention mechanism layer Multi-headAttention and the fully connected feedforward network FFN of each encoder Transformer Encoder, a genetic algorithm is used to search for the optimal pruning strategy; The optimal pruning strategy searched is used for pruning operation; Wherein, for the multi-head attention mechanism layer Multi-headAttention and the fully connected feedforward network FFN of each encoder Transformer Encoder, a genetic algorithm is used to search for the optimal pruning strategy, including: For different pruning rate configurations p i = [r1, r2,..., rn], where r1, r2,..., rn represent the pruning ratio for the 1st to nth Multi-head Attention layers, and s1, s2,..., sn represent the pruning ratio for the 1st to nth Fully Connected Feed-Forward Networks (FFNs). i n n n n The search target is as follows:​​​​ Obj(p i ) = aAccuracy(p i ) + bAFLOPs(p i ); (7) Wherein, TP is the number of samples correctly predicted as positive class, TN is the number of samples correctly predicted as negative class, FP is the number of negative class samples incorrectly predicted as positive class, and FN is the number of positive class samples incorrectly predicted as negative class; wherein FLOPs ori denotes the number of floating point operations of the original model, denotes the number of floating point operations of the pruned model at pruning rate configuration p i after pruning. Accuracy(p i ) represents the accuracy of the model on the test set under pruning strategy p i , and ΔFLOPs(p i ) represents the difference between the FLOPs of the model obtained by the special computing module and the target FLOPs under pruning strategy p i , and the target is to maximize the objective function Obj(p i ). The special calculation module uses the following formula: where d model is the fixed dimension of each patch projected by a linear layer, N is the number of patches obtained by segmentation, d patch is the dimension of each patch, and L is the number of encoder Transformer Encoders. α and β represent the proportion of accuracy and FLOPs difference respectively when evaluating different pruning rate configurations, and the search space is defined as follows: O∈(0.001，0.5)(8) O represents that the pruning rate is searched for the optimal value between 0.001 and 0.5; Randomly initialize a set of candidate pruning strategies within the range of the search space, each candidate pruning strategy as an individual represents a set of candidate solutions, combine all candidate solutions into a population set, and each population set is updated through iteration. The candidate pruning strategy is represented as follows: Pol FFN = {r1, r2,..., r 11 , r 12} (9) Pol head = {s1, s2,..., s 11 , s 12} (10) Pol represents the result of the candidate pruning strategy, and the right side of equation (9) {r1, r2, ..., r 11 r 12 } represents the pruning ratio of the fully connected feedforward network (FFN) in the Transformer Encoder from layers 1 to 12, denoted as Pol. FFN The right side of equation (10) [s1, s2, ..., s 11 s 12 ] represents the pruning ratio of the multi-head attention mechanism layers in the Transformer Encoder from layers 1 to 12, denoted as Pol. head ; The fitness function is set the same as the objective function, and the expression is as follows: Obj(p i ) = a Accuracy(p i ) + b AFLOPs(p i ); After defining the fitness function, the initial candidate pruning strategy evolves into a new generation of population through selection, crossover and mutation operations; the evolution process adopts the roulette wheel selection mechanism to completely replace the original population; after 50 iterations, the algorithm converges to the optimal pruning strategy. 7.The method of claim 6, wherein the pruning the deep learning model comprises: determining a pruning rate based on the number of the pruning operations; and pruning the deep learning model based on the pruning rate. The optimal pruning strategy searched is used for pruning operations, including: Pruning method of a fully connected feed-forward network FFN: the fully connected feed-forward network FFN is known to comprise an input layer, an intermediate layer and an output layer connected in sequence; at the intermediate layer of the fully connected feed-forward network FFN, the number of neurons is reduced by pruning, for the fully connected feed-forward network FFN of the layer Transformer Encoder, the feature Z l , the feature of each neuron j is the intermediate layer is implemented by a hidden layer; By computing the Gram matrix and dependency scores evaluating an importance score for each neuron; gram matrix the elements in the gram matrix G are defined as where k is a kernel function that computes the similarity between two vectors; Dependency score The formula is: wherein, is the Gram matrix defined on tr denotes the trace of a matrix, CLC is a constant matrix; the dependency score is the importance score of each neuron; According to the neuron pruning ratio k searched by the genetic algorithm l , the neurons with the highest importance scores are retained; the rest are deleted ; Pruning method of Multi-head Attention layer: In the multi-head attention mechanism layer Multi-head Attention of the first layer Transformer Encoder, the output feature of each attention head h is represented as computing a gram matrix K for each attention head l and dependency scores evaluating an importance score for each head; According to the search, the pruning ratio σ l The top σ l ·N h attention heads are retained, and the rest are deleted, N h is the total number of attention heads.

8. An image classification system based on a pruned deep learning model, characterized in that, including: The pruning module is configured to: obtain a deep learning model, perform pruning operation on the deep learning model to obtain a pruned deep learning model; wherein the deep learning model is a Vision Transformer model; the pruning operation includes: using a genetic algorithm to perform pruning operation on the Multi-head Attention layer and the FFN of each Transformer Encoder in the Vision Transformer model; the pruning operation further includes: adding a Token Select layer between the Multi-head Attention layer and the FFN of each Transformer Encoder, the Token Select layer calculates an importance score for each token, and completes the pruning operation according to the importance score to reduce the number of tokens; The training module is configured to: construct a training set, the training set being an image with known image classification label; input the training set into the deep learning model to train the model and obtain a trained deep learning model; copy the model parameters of the trained deep learning model to the pruned deep learning model to obtain a trained pruned deep learning model; The classification module is configured to: obtain an image to be classified, input the image to be classified into the trained pruned deep learning model to obtain a classification result of the image.

9. An electronic device, comprising: a memory for non-transiently storing computer readable instructions; and a processor for running the computer readable instructions, wherein the computer readable instructions, when executed by the processor, perform the method of any one of claims 1-7.

10. A storage medium characterized by, non-transitory computer readable instructions, wherein when the non-transitory computer readable instructions are executed by a computer, the method of any one of claims 1-7 is performed.

Images