Image classification method and system based on pruning deep learning model

Abstract

The invention discloses an image classification method and system based on a pruning deep learning model, and the method comprises the steps: obtaining a deep learning model, carrying out the pruning operation of the deep learning model, and obtaining a pruning deep learning model; a pruning operation is carried out, and the pruning operation comprises the step of carrying out pruning operation on a multi-head attention mechanism layer (Multi-head attention) and a full-connection feed-forward network (FFN) of each encoder Transform Ender by adopting a genetic algorithm in a Vision Transform model, and the step of carrying out pruning operation on the multi-head attention mechanism layer (Multi-head attention) and the full-connection feed-forward network (FFN) of each encoder Transform Ender; obtaining a trained pruning deep learning model; and obtaining a to-be-classified image, and inputting the to-be-classified image into the trained pruning deep learning model to obtain an image classification result.

Description

Technical Field

[0001] The present invention relates to the technical field of deep learning model pruning and compression, and particularly to an image classification method and system based on a pruned deep learning model. Background Art

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

[0003] With the continuous improvement of model architectures, in the field of natural language processing, the Transformer model has made revolutionary progress. With the continuous evolution of this model architecture, a large number of recent studies have begun to explore the application of Vision Transformer (ViT) based on Transformer in computer vision tasks, especially large-scale image classification. Compared with traditional convolutional neural networks (CNNs), ViT has shown significant performance advantages in processing large-scale image classification tasks, and its entire structure does not contain any convolutional operations. DeiT further optimizes the structure of ViT by introducing a knowledge distillation strategy and achieves performance beyond the original ViT when only using the ImageNet-1k dataset. In addition, DeepViT proposes a simple and effective Re-attention mechanism to further improve the model performance by reconstructing the attention map. With the continuous development of ViT and its derivative models, their application scope has been extended to multiple task areas related to images, including image segmentation, object detection, etc., showing the wide applicability and potential of the ViT architecture in the field of computer vision.

[0004] Although Vision Transformer (ViT) performs excellently in large-scale classification tasks, it still has the problem of high complexity, mainly stemming from its processing flow and attention mechanism: First, ViT divides the input image into multiple small patches; then, these patches are converted into tokens, and each token corresponds to an element in the input sequence. Since ViT uses the self-attention mechanism to process these tokens, it needs to calculate the relationship between each token and all other tokens in the sequence, which involves a large number of parameters and calculations. Especially when the number of tokens increases, the computational complexity will increase quadratically. In addition, the depth of the ViT model also increases the computational burden because 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 a large model size and high computational costs. Further compressing the model to reduce the resources required during runtime and enabling wider applications has become a practical need. Summary of the Invention

[0005] To solve the problems of excessive number of parameters and large consumption of computing resources in the Vision Transformer model, the present invention proposes an image classification method and system based on pruning a deep learning model, and this pruning framework is denoted as MSCViT. The present invention reduces the computational amount and the number of parameters of the model to a certain extent by pruning three modules in the model, namely tokens, attention heads, and FFNs. The importance evaluation method for the token module is improved so that it does not require a large amount of additional computing resources, and a hierarchical pruning method based on intelligent optimization algorithms is adopted for the attention head and FFN modules. Through the above methods, the number of parameters and FLOPs of the model are effectively reduced, and the accuracy is maintained unchanged or only slightly decreased.

[0006] On the one hand, an image classification method based on pruning a deep learning model is provided, including:

[0007] Obtain a deep learning model, 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 operations include: performing pruning operations on the multi-head attention mechanism layer Multi-head Attention and the fully connected feed-forward network FFN of each encoder Transformer Encoder in the Vision Transformer model using a genetic algorithm; the pruning operations also include: adding a Token Select layer between the multi-head attention mechanism layer Multi-head Attention and the fully connected feed-forward network FFN of each encoder Transformer Encoder, and the Token Select layer calculates the importance score for each token and completes the pruning operation according to the importance score to reduce the number of tokens;

[0008] Construct a training set, where the training set is an image with a known image classification label; input the training set into the deep learning model to train the model to obtain a trained deep learning model; copy the model parameters of the trained deep learning model into the pruned deep learning model to obtain a trained pruned deep learning model;

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

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

[0011] A pruning module, which is configured to: obtain a deep learning model, 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 operations include: performing pruning operations on the multi-head attention mechanism layer Multi-head Attention and the fully connected feed-forward network FFN of each encoder Transformer Encoder in the Vision Transformer model using a genetic algorithm; the pruning operations also include: adding a Token Select layer between the multi-head attention mechanism layer Multi-headAttention and the fully connected feed-forward network FFN of each encoder Transformer Encoder, and the Token Select layer calculates the importance score for each token and completes the pruning operation according to the importance score to reduce the number of tokens;

[0012] A training module, which is configured to: construct a training set, where the training set is images with known image classification labels; input the training set into a deep learning model to 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.

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

[0014] In another aspect, an electronic device is further provided, including:

[0015] A memory for non-temporarily storing computer-readable instructions; and

[0016] A processor for running the computer-readable instructions,

[0017] wherein when the computer-readable instructions are run by the processor, the method described in the first aspect above is executed.

[0018] In another aspect, a storage medium is further provided, which non-temporarily stores computer-readable instructions, wherein when the non-temporary computer-readable instructions are executed by a computer, the method described in the first aspect is executed.

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

[0020] The above technical solutions have the following advantages or beneficial effects:

[0021] (1) Pruning Tokens can streamline the core structure of the Vision Transformer (ViT) model, that is, reduce the number of Tokens participating in self-attention calculation, thereby reducing the model's parameter quantity and computational complexity. Since ViT divides the input image into multiple Tokens, each Token representing a small patch of the image, pruning can remove those Tokens that contribute less 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. At the same time, the pruned model may be more stable during the training process due to fewer parameters, helping the model learn more robust feature representations. In addition, pruning enhances the interpretability of the model because fewer but more important Tokens are retained, enabling researchers to more intuitively understand the basis for the model's decisions, and thus more confidently deploy and use these optimized models in practical applications.

[0022] (2) The use of genetic algorithms to search for the optimal pruning configurations of different layers of the Attention head and FFN modules proposed in this invention shows significant beneficial effects. Pruning different layers of the neural network with specifically optimized different pruning rates can precisely adjust the sparsity of the model and optimize performance and efficiency. This method makes full use of the characteristic that the importance of each layer in feature extraction is different, allowing more parameters to be retained for key layers to capture important information, while more aggressive pruning is performed on less important layers to reduce redundancy. This differential treatment not only improves the model's generalization ability and avoids overfitting but also balances the use of computing resources, enabling the model to reduce the inference time and memory occupancy while maintaining accuracy, enhancing the practicality and flexibility of the model in resource-constrained environments.

[0023] (3) The weight sharing mechanism in this invention greatly reduces the time required for retraining compared to retraining each compressed model and significantly improves the search efficiency of the genetic algorithm. BRIEF DESCRIPTION OF THE DRAWINGS

[0024] The specification drawings forming a part of this invention are used to provide a further understanding of this invention. The schematic embodiments of this invention and their descriptions are used to explain this invention and do not constitute an improper limitation of this invention.

[0025] Figure 1 It is a structural diagram of the Vision Transformer model;

[0026] Figure 2 It is a schematic diagram of the attention weight analysis method of this invention;

[0027] Figure 3Schematic diagram of the token selection layer of the present invention;

[0028] Figure 4 Schematic diagram of token fusion pruning of the present invention;

[0029] Figure 5 Schematic diagram of the genetic algorithm search pruning configuration of the present invention;

[0030] Figure 6 Schematic diagram of the internal connection relationship of the fully connected feed-forward network FFN of the present invention. Detailed implementation manners

[0031] It should be noted that the following detailed description is exemplary and is intended to provide further illustration of the present invention. Unless otherwise specified, all technical and scientific terms used herein have the same meaning as commonly understood by those of ordinary skill in the technical field to which the present invention belongs.

[0032] Embodiment 1

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

[0034] The image classification method based on a pruned deep learning model includes:

[0035] S101: Obtain a deep learning model, perform a 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: performing a pruning operation on the multi-head attention mechanism layer Multi-head Attention and the fully connected feed-forward 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 Multi-head Attention and the fully connected feed-forward network FFN of each encoder Transformer Encoder, and 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;

[0036] S102: Construct a training set, where the training set is an image with a known image classification label; input the training set into the deep learning model, train the model to 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;

[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] Further, as Figure 1 shown, the Vision Transformer model includes:

[0039] A flattened layer (Linear Projection of Flattened Patches), a position encoding layer, a series of L encoder Transformer Encoders, a multi-layer perceptron, and a classification layer connected in sequence.

[0040] Further, the Vision Transformer model includes:

[0041] By dividing the input image into blocks;

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

[0043] The position encoding layer performs a linear projection on the one-dimensional vector and adds position encoding to obtain a number of tokens;

[0044] The series of L encoder Transformer Encoders perform L times of feature extraction on the tokens to obtain the extracted features;

[0045] Input the extracted features into the multi-layer perceptron for calculation, and the classification layer gives the classification label according to the calculation result of the multi-layer perceptron.

[0046] Exemplarily, the flattened layer divides the input image into image patches of a fixed size, then flattens each image patch into a one-dimensional vector, and then maps the vector to a high-dimensional embedding space through a linear projection (convolutional layer).

[0047] Exemplarily, the position encoding layer generates the position encoding PE by using sine and cosine functions, and the formula is as follows:

[0048]

[0049]

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

[0051] Among them, there are L encoder Transformers Encoder connected in series in the Vision Transformer model. Each encoder Transformer Encoder includes:

[0052] A first-layer normalization module LayerNorm, a multi-head attention mechanism layer Multi-headAttention, a first adder, a second-layer normalization module LayerNorm, a fully-connected feed-forward network FFN, and a second adder connected in sequence;

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

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

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

[0056] Exemplarily, 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 encoder Transformers Encoder connected in series, a multi-layer perceptron, and a classification layer connected in sequence.

[0058] Among them, the difference between the pruned deep learning model and the Vision Transformer model is that:

[0059] A Token Select layer is added between the multi-head attention mechanism layer Multi-head Attention and the fully-connected feed-forward network FFN of each encoder Transformer Encoder of the pruned deep learning model; the input end of the Token Select layer is connected to the output end of the multi-head attention mechanism layer Multi-head Attention, and the output end of the Token Select layer is connected to the input end of the fully-connected feed-forward network FFN;

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

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

[0062] For each encoder Transformer Encoder of the pruned deep learning model, the Multi-head Attention layer and the fully connected feed-forward network FFN both use the genetic algorithm for pruning operations.

[0063] Furthermore, the Token Select layer calculates the importance score for each token and completes the pruning operation according to 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, fuse the tokens with scores lower than the threshold, and retain the tokens with scores higher than the threshold.

[0066] Furthermore, the (1-1) calculating the importance score for each token specifically includes:

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

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

[0069]

[0070] Among them, Q h , V h are respectively the query matrix, key matrix, and value matrix converted from the input tokens in each head h, and d k is the scaling factor;

[0071] Derive the attention distribution matrix of a single token to other tokens:

[0072]

[0073] Among them, A h [i, j] represents the attention weight of the i-th token to the j-th token, represents the query value of the i-th token in the h-th attention head, represents the key value of the j-th token in the h-th attention head, 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: use the query vector of the i-th token and the key vector of each token to do a dot product, after scaling and softmax normalization, a probability distribution is obtained, and the probability distribution is the attention distribution matrix A h [i, j]; the probability distribution represents the attention distribution of the i-th token to all other tokens.

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

[0076]

[0077] Among them, 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 degree 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 the role in feature extraction;

[0078] Since the model is based on the multi-head self-attention mechanism, it means that there are multiple attention heads inside each encoder Transformer Encoder. 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:

[0079]

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

[0081] The importance score of each token is expressed as:

[0082]

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

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

[0085] Furthermore, the step (1-2) fuses the tokens with scores lower than the score threshold, and the specific fusion process includes:

[0086] (1-2-1) According to the importance scores of the tokens, identify several tokens with importance scores lower than the set threshold;

[0087] (1-2-2) Perform a weighted summation operation on several tokens with importance scores lower than the threshold to generate a fused token;

[0088] The specific formula is as follows:

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

[0090] Among them, x fusedThe fused token, N is a set of several tokens whose importance scores are lower than the threshold, a i is the weight of the token corresponding to index i, x i is the token at index i.

[0091] (1 - 2 - 3) Add the calculated fused token to the set of tokens whose importance scores are higher than the set threshold.

[0092] It should be understood that after the calculation of the token importance in the selected encoder Transformer Encoder of the present invention, in order to implement token pruning, a Token Selection Layer (T - Select) is added between the MHSA and FFN modules in each Transformer layer, as Figure 3 shown. This operation does not add additional parameters. The Token Selection Layer receives the input feature Z l , and then according to the given pruning ratio v 1 selects and retains the most important (1 - v l ) N tokens, and fuses the least important v 1 N tokens through importance score fusion. These scores are calculated using the attention weight analysis method. After fusing the tokens with scores lower than the score threshold, the token length of the subsequent layer will be reduced, and the subsequent Token Selection Layer will select from the tokens retained by the previous layer, as Figure 4 shown.

[0093] Furthermore, in the Vision Transformer model, pruning operations are performed on the multi - head attention mechanism layer Multi - head Attention and the fully - connected feed - forward network FFN of each encoder Transformer Encoder, including:

[0094] (2 - 1): Use the genetic algorithm to search for the optimal pruning strategy for the multi - head attention mechanism layer Multi - headAttention and the fully - connected feed - forward network FFN of each encoder Transformer Encoder;

[0095] (2 - 2) Perform pruning operations using the searched optimal pruning strategy.

[0096] Furthermore, as Figure 5 shown, the (2 - 1): Use the genetic algorithm to search for the optimal pruning strategy for the multi - head attention mechanism layer Multi - head Attention and the fully - connected feed - forward network FFN of each encoder Transformer Encoder, including:

[0097] (2-1-1) For different pruning rate configurations p i = [r 1 , r 2 ,..., r n , s 1 , s 2 ,..., s n , where r 1 , r 2 ,..., r n represents the pruning ratio for the first to the n-th multi-head attention mechanism layers Multi-head Attention, and s 1 , s 2 ,..., s n represents the pruning ratio for the first to the n-th fully connected feed-forward network FFN. 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 class, TN is the number of samples correctly predicted as negative class, FP is the number of negative class samples wrongly predicted as positive class, and FN is the number of positive class samples wrongly predicted as negative class.

[0101]

[0102] Where, represents the number of floating-point operations of the original model, represents the number of floating-point operations of the pruned model under the pruning rate configuration p i .

[0103] Accuracy(p i ) represents the accuracy of the model on the test set under the pruning strategy p i , and ΔFLOPs(p i ) represents the difference between the FLOPs of the model obtained by the dedicated calculation module and the target FLOPs under the pruning strategy p i . The goal is to maximize the objective function Obj(p i );

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

[0105]

[0106] Among them, d model is the fixed dimension projected by each patch through 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.

[0107] α and β respectively represent the proportions of the accuracy and FLOPs difference when evaluating different pruning rate configurations. The larger α is, the higher the proportion of the accuracy factor. If β is larger, it means the higher proportion of the 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 searched for the optimal value between 0.001 and 0.5. Through actual analysis, if the pruning ratio is too high, it will cause a large drop in accuracy. Therefore, the pruning ratio is restricted between 0.001 and 0.5 to achieve a balance between accuracy and model calculation amount.

[0111] (2-1-3) Within the range of the search space, randomly initialize a set of candidate pruning strategies. Each candidate pruning strategy is regarded as an individual, representing a set of candidate solutions. Combine all candidate solutions into a population set, and each population set undergoes iterative strategy updates. The candidate pruning strategy is expressed as follows:

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

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

[0114] Pol represents the result of the candidate pruning strategy. The right side of equation (9) {r 1 , r 2 ,..., r 11 , r 12}} represents the pruning ratio of the fully connected feed-forward network FFN in the 1st to 12th layer encoder Transformer Encoder, denoted as Pol FFN , and the right side of equation (10) [s 1 , s 2,...,s 11 ,s 12 represents the pruning ratio of the multi-head attention mechanism layer in the 1st to 12th layer encoder Transformer Encoder, denoted as Pol head ;

[0115] (2-1-4) The fitness function is set the same 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 roulette wheel selection mechanism is adopted in the evolution process, and the offspring completely replace the original population. After 50 iterations, the algorithm converges to the optimal pruning strategy.

[0118] Furthermore, the pruning operation using the optimal pruning strategy found includes:

[0119] Pruning method for the fully connected feedforward network FFN: As Figure 6 shown, it is known that the fully connected feedforward network FFN includes an input layer, an intermediate layer, and an output layer connected in sequence; in the intermediate layer of the fully connected feedforward network FFN, the number of neurons is reduced through pruning. For the fully connected feedforward network FFN of the l-th layer Transformer Encoder, there is a feature Z l , and the feature of each neuron j is The intermediate layer is implemented through a hidden layer;

[0120] By calculating the Gram matrix and the dependence score to evaluate the importance score of each neuron;

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

[0122] The dependence score has the formula:

[0123]

[0124] where, is defined on The Gram matrix on it, tr represents the trace of the matrix, and CLC is a constant matrix; the dependence score, which is the importance score of each neuron;

[0125] The neuron pruning ratio k obtained by searching with the genetic algorithm l , retain the neurons with the highest importance scores neurons; the rest are deleted.

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

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

[0128] Calculate the Gram matrix K of each attention head l and the dependence score Evaluate the importance score of each head;

[0129] According to the pruning ratio σ obtained by searching l , the top σ l ·N h attention heads with the highest scores are retained, and the rest are deleted, where N h is the total number of attention heads.

[0130] It should be understood that using the genetic algorithm to search for different pruning ratios for each layer for a single attention head and FFN of the multi-head attention mechanism layer Multi-head Attention. Through such targeted pruning, a better balance can be found between reducing the computational burden and maintaining the model performance. It allows the model to maintain high performance at critical layers while reducing the use of computational resources at layers that have less impact on the final output.

[0131] In Vision Transformer, a single attention head and FFN of the multi-head attention mechanism layer Multi-head Attention are also important components and functional parts. For these two modules, the genetic algorithm is used to find the pruning strategy for each layer in FFN and a single attention head of the multi-head attention mechanism layer Multi-head Attention, so that the ViT network compressed by FFN and a single attention head of the multi-head attention mechanism layer Multi-head Attention can find the optimal solution by weighing the computational amount, accuracy, and inference time under the pruning range constraint.

[0132] Copy the model parameters of the trained deep learning model to the pruned deep learning model. The present invention uses an improved model accuracy evaluation method with shared weights 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 under each pruning strategy, we efficiently evaluate the accuracy of the model through weight sharing. We pre-train the complete model weights W so that we can predict the accuracy after independent training of the sub-model based on the weights inherited by the sub-model. Thanks to this sharing mechanism, we only need to train the model once to obtain the corresponding weights and then use these weights to evaluate the accuracy of multiple pruning strategies. This method greatly reduces the time required compared to retraining each compressed model and significantly improves the search efficiency of the genetic algorithm here.

[0134] This embodiment provides a compression and retraining method based on a modular compression framework, which is divided into two stages as a whole: 1) Pruning stage. It involves removing unimportant weights or neurons in the ViT model. For the token module, the importance of tokens 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 the dependence score, and hierarchical pruning is performed according to the searched pruning rate. The purpose of this stage is to reduce the complexity and computational cost of the model while trying to retain the key information of the model. 2) The retraining stage is carried out after pruning. The sub-dataset ILSVRC-2012 of the ImageNet dataset is used, and the remaining weights are fine-tuned using the AdamW optimizer and the cosine decay strategy to recover part of the performance loss that may be caused by pruning, ensuring that the model still maintains a high accuracy and generalization ability while reducing the number of parameters.

[0135] First, calculate the importance score of each token by averaging the attention weights of each token in the selected encoder Transformer Encoder to all tokens (including itself). After the calculation of the token importance in the selected encoder Transformer Encoder is completed, in order to implement token pruning, a Token selection layer (T-Select) is added between the MHSA and FFN modules in each Transformer layer, and this operation does not increase additional parameters. T-Select receives the input feature Z l and then selects the most important (1 - v l ) N tokens according to the given pruning ratio v l and fuses the least important v lN tokens, and these scores are calculated using the attention weight analysis method. Finally, the tokens with scores below the score threshold are fused, the length of the tokens in the subsequent layers will be reduced, and the subsequent T-Select will select from the tokens retained in the previous layer.

[0136] For a single attention head and the FFN module of the multi-head attention mechanism layer Multi-head Attention, first use the genetic algorithm to search for the optimal pruning strategy for different layers, randomly initialize a set of pruning strategies within the search space, and set the fitness function and the objective function. The initial candidate pruning strategies evolve to generate a new generation of populations through selection, crossover, and mutation operations. This process adopts the roulette wheel selection mechanism to completely replace the original population with the offspring. After 50 iterations, the algorithm converges to an approximate optimal pruning strategy. In the middle layer of the FFN module, according to the neuron pruning ratio k l , retain the most important neurons. Evaluate the importance of each neuron by calculating the Gram matrix and the dependency score . The neurons with the highest scores are retained, and the rest are pruned. For a single attention head of the multi-head attention mechanism layer Multi-head Attention, in the multi-head self-attention (MHSA) of the l-th layer, the output feature of each head h is represented as Z l . According to the pruning ratio σ l obtained by searching, simplify the output feature of each head through average pooling to form Then calculate the Gram matrix K l of each head and the dependency score to evaluate the importance of each neuron. The neurons with the highest scores are retained, and the rest are pruned.

[0137] Fine-tune the pruned model to restore the accuracy. Use the Pytorch framework to train on a specific GPU. The configuration used for training follows the initial configuration of the ViT model, adopts the AdamW optimizer and the cosine decay strategy, aiming to achieve a more stable learning rate adjustment during the training process. At the same time, reduce the learning rate and the batch size to more finely control the convergence process of the model and ensure that the model can achieve better generalization ability during training. After the model is initially converged, perform 100 / 80 times of fine-tuning on the pruned ViT model respectively to further improve the performance of the model.

[0138] The pruning method for the tokens module uses the method of attention weight analysis to evaluate the importance of each token, and prunes and removes the unimportant tokens according to the importance ranking to reduce the burden of the model in this part.

[0139] The pruning method of a single attention head of the multi-head attention mechanism layer Multi-head Attention uses an intelligent optimization algorithm to search for the optimal pruning strategy for each layer in the model to achieve accurate pruning of 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 to achieve pruning effects with different pruning rates for different layers.

[0141] The model performance is restored by fine-tuning and retraining, and the performance of the model is restored to the level before pruning as much as possible.

[0142] The attention weight analysis method utilizes the characteristics of the attention mechanism itself and calculates the attention size between a single token and other tokens as the importance judgment criterion.

[0143] The pruning of unimportant tokens does not adopt the direct removal method, but merges them into one token based on their characteristics to reduce 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 layer of FFN modules 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 mechanism layer, a genetic algorithm is used to search for the optimal pruning ratio of a single attention head in each layer of the Multi-head Attention mechanism layer in the ViT model to adapt to the number distribution of single attention heads in the Multi-head Attention mechanism layer in different levels of the model.

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

[0147] The AdamW optimizer and cosine decay strategy are used to achieve more stable learning rate adjustment during training. At the same time, the learning rate and batch size are reduced to more finely control the convergence process of the model.

[0148] The present invention discloses a modular compression method for Vision Transformer, which relates to the field of machine learning and aims to reduce the computational amount and the number of parameters required by the model through pruning to adapt to resource-constrained environments. This method is achieved through two steps: pruning the ViT model and fine-tuning the pruned model to restore performance. Pruning the ViT model involves three modules: tokens, individual attention heads of the multi-head attention mechanism layer (Multi-head Attention), and FFN. For the token module, an attention weight analysis method is used to evaluate the importance of tokens, and a token fusion method is adopted to reduce the number of tokens. For the individual attention heads of the multi-head attention mechanism layer (Multi-head Attention) and the FFN module, a genetic algorithm is used to search for the optimal pruning configuration for each layer. During the search process, a weight sharing method is also adopted to rapidly improve the search speed, and the importance of each neuron is evaluated based on the Gram matrix and the dependence score. Neurons below the threshold are pruned. When fine-tuning the pruned model, the AdamW optimizer and the cosine decay strategy, as well as a low learning rate and batch size, are used to better control the convergence process of the model and restore the performance of the model to the level before pruning as much as possible. The technical solution of the present invention reduces the computational amount and the number of parameters of the Vision Transformer model, reduces the computational resources required during the model training process, speeds up the inference speed of the model, makes the ViT model easier to be deployed to different platforms, and reduces the computational and storage costs.

[0149] Embodiment 2

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

[0151] A pruning module, which is configured to: obtain a deep learning model, perform a 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: performing a pruning operation on the multi-head attention mechanism layer (Multi-head Attention) and the fully connected feed-forward network (FFN) of each encoder (TransformerEncoder) in the Vision Transformer model by using a genetic algorithm; the pruning operation further includes: adding a Token Select layer between the multi-head attention mechanism layer (Multi-headAttention) and the fully connected feed-forward network (FFN) of each encoder (Transformer Encoder), and 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.

[0152] A training module, which is configured to: construct a training set, where the training set is an image with a known image classification label; input the training set into a deep learning model to 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.

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

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

[0155] In the above embodiments, the descriptions of each embodiment have their own emphases. For parts not detailed in a certain embodiment, reference can be made to the relevant descriptions of other embodiments.

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

[0157] Embodiment 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 above one or more computer programs are stored in the memory. When the electronic device runs, the processor executes the one or more computer programs stored in the memory so that the electronic device executes the method described in Embodiment 1 above.

[0159] It should be understood that in this embodiment, the processor can be a central processing unit CPU, and the processor can also be other general-purpose processors, digital signal processors DSP, application-specific integrated circuits ASIC, off-the-shelf programmable gate arrays FPGA, or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. The general-purpose processor can be a microprocessor or the processor can also be any conventional processor, etc.

[0160] The memory may include a read-only memory and a random access memory, and provide instructions and data to the processor. A part of the memory may also include a non-volatile random access memory. For example, the 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 circuit in the hardware of the processor or the instructions in the form of software.

[0162] The method in the first embodiment can be directly embodied as being executed by the hardware processor, or executed by the combination of the hardware and software modules in the processor. The software module can be located in the mature storage media in the art such as random access memory, flash memory, read-only memory, programmable read-only memory, or electrically erasable programmable memory, registers, etc. This storage media is located in the memory, and the processor reads the information in the memory and combines its hardware to complete the steps of the above method. To avoid repetition, it will not be described in detail here.

[0163] Those of ordinary skill in the art can realize that the units and algorithm steps of each example described in combination with this embodiment can be implemented by electronic hardware or the combination of computer software and electronic hardware. Whether these functions are executed in a hardware or software manner depends on the specific application and design constraints of the technical solution. Professional technicians can use different methods to implement the described functions for each specific application, but this implementation should not be considered to exceed the scope of the present invention.

[0164] Embodiment 4

[0165] This embodiment also provides a computer-readable storage medium for storing computer instructions. When the computer instructions are executed by the processor, the method described in the first embodiment is completed.

[0166] The above are only the preferred embodiments of the present invention and are not used to limit the present invention. For those skilled in the art, the present invention can have various changes and modifications. Any modification, equivalent replacement, improvement, etc. made within the spirit and principle of the present invention shall be included within the protection scope of the present invention.

Claims

1. Image classification method based on pruning deep learning model, characterized by: include: Acquire a deep learning model, perform a pruning operation 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: using a genetic algorithm to perform a pruning operation on 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; the pruning operation also includes: 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 completes the pruning operation according to the importance score to reduce the number of tokens; Constructing a training set, wherein the training set is an image with a known image classification label; inputting the training set into a deep learning model, training the model, and obtaining a trained deep learning model; copying the model parameters of the trained deep learning model into a pruned deep learning model, and obtaining a trained pruned deep learning model; Obtain the image to be classified, input the image to be classified into the trained pruned deep learning model, and obtain the classification result of the image.

2. The image classification method based on the pruned deep learning model as claimed in claim 1, characterized in that: The pruned deep learning model includes: a flattening layer, a position encoding layer, a series of L encoders, a multi-layer perceptron, and a classification layer; The difference between the pruning deep learning model and the Vision Transformer model is: Add a Token Select layer between the Multi-head Attention layer and the fully connected feedforward network FFN of each encoder Transformer Encoder of the pruned deep learning model; 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; The difference between the pruned deep learning model and the Vision Transformer model is 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 pruned using genetic algorithms.

3. The image classification method based on the pruned deep learning model as claimed in claim 1, characterized in that: The TokenSelect layer calculates the importance score for each token and performs pruning operations based on the importance score to reduce the number of tokens, including: (1-1) Calculate the importance score of each token; (1-2) Compare the importance score of each token with the set score threshold, merge the tokens below the score threshold, and retain the tokens above the score threshold.

4. The image classification method based on the pruned deep learning model as claimed in claim 3, characterized in that: Calculate the importance score of each token, including: In Vision Transformer, the input sequence X is mapped to three matrices: query, key, and value. In each encoder, the attention weight A of each head in the multi-head attention mechanism is calculated as shown in formula (2): Among them, Q h , V h The query matrix, key matrix and value matrix converted from the token token input in each header h, respectively, d k is the scaling factor; The attention distribution matrix of a single token to other tokens is calculated: Among them, A h [i,j] represents the attention weight of the i-th token to the j-th token. represents the query value of the i-th token token in the h-th attention head, represents the key value of the jth token in the hth attention head, Represents the key value of the mth token in the hth attention head; Within a single attention head, the importance score of the i-th token is calculated by averaging the attention weights of the i-th token over all tokens: in, represents the initial importance score of the i-th token in the h-th attention head; Aggregate and average the initial importance scores of the same token in all attention heads in the encoder Transformer Encoder to get the final importance score of the i-th token in the entire encoder Transformer Encoder: in, Represents the final importance score of the i-th token token in the current encoder Transformer Encoder; the importance score of each token token is expressed as: Among them, Imp i Represents the importance score of the i-th token token.

5. The image classification method based on the pruned deep learning model as claimed in claim 4, characterized in that: Tokens below the score threshold are merged. The specific fusion process includes: (1-2-1) According to the importance score of the token, identify several tokens whose importance scores are lower than the set threshold; (1-2-2) Perform a weighted sum operation on several tokens whose importance scores are lower than the threshold to generate a fused token; The specific formula is as follows: x fused =∑ i∈N a i x i Among them, x fused is the fused token, N is a set of tokens whose importance scores are lower than the threshold, a i is the weight of the token corresponding to index i, x i is the token at index i; (1-2-3) Add the calculated fused token to the token set whose importance score is higher than the set threshold.

6. The image classification method based on the pruned deep learning model as claimed in claim 1, characterized in that: In the VisionTransformer model, the genetic algorithm is used to prune the multi-head attention mechanism layer Multi-headAttention and the fully connected feedforward network FFN of each encoder Transformer Encoder, including: A genetic algorithm is used to search for the optimal pruning strategy for each encoder Transformer Encoder's multi-head attention mechanism layer and the fully connected feedforward network FFN; Use the searched optimal pruning strategy to perform pruning operations; Among them, the genetic algorithm is used to search for the optimal pruning strategy for each encoder Transformer Encoder's multi-head attention mechanism layer Multi-headAttention and fully connected feedforward network FFN, including: For different pruning rate configurations p i =[r1,r2,...,r i ,s n ,s2,...,s n ], where r1, r2, ..., r n Indicates the pruning ratio of the 1st to the nth multi-head attention mechanism layer Multi-head Attention, s1,s2,...,s n Indicates the pruning ratio of the 1st to the nth fully connected feedforward network FFN. The search target is as follows: Obj(p i )=αAccuracy(p i )+βΔFLOPs(p i ); (7) Among them, 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. Among them, FLOPs ori Indicates the number of floating point operations of the original model, Indicates that the pruning rate configuration p i The number of floating point operations of the pruned model; 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 FLOPs of the model obtained by the dedicated computing module and the target FLOPs is to make the objective function Obj(p i )maximize; The dedicated calculation module uses the following formula: Among them, d model is the fixed dimension of each patch projected through a linear layer, N is the number of patches obtained by segmentation, and d patch is the dimension of each patch, L is the number of encoder Transformer Encoder; α and β represent the proportion of accuracy and FLOPs difference when evaluating different pruning rate configurations, respectively. The search space is defined as follows: O∈(0.001，0.5)(8) O means that the pruning ratio is to find the optimal value between 0.001 and 0.5; In the search space, a set of candidate pruning strategies are randomly initialized. Each candidate pruning strategy is regarded as an individual, representing a set of candidate solutions. All candidate solutions are combined into a population set. Each population set is updated through iteration. The candidate pruning strategies are expressed as follows: Pol FFN ={r1,r2,...,r 11 ,r 12 } (9) Pol head ={s1,s2,...,s 11 ,s 12 } (10) Pol represents the candidate pruning strategy result. 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 1st to 12th layer encoder Transformer Encoder, denoted by Pol FFN , the right side of equation (10) [s1, s2, ..., s 11 ,s 12 ] represents the pruning ratio of the Multi-head Attention layer in the 1st to 12th layer encoder Transformer Encoder, denoted by Pol head ; The setting of the fitness function is the same as the objective function, and the expression is as follows: Obj(p i )=αAccuracy(p i )+βΔFLOPs(p i ); After defining the fitness function, the initial candidate pruning strategy evolves to produce a new generation of population through selection, crossover and mutation operations; the evolution process adopts a roulette selection mechanism to completely replace the original population with the offspring; after 50 iterations, the algorithm converges to the optimal pruning strategy.

7. The image classification method based on the pruned deep learning model as claimed in claim 6, characterized in that: Use the searched optimal pruning strategy to perform pruning operations, including: Pruning method of fully connected feedforward network FFN: It is known that the fully connected feedforward network FFN includes an input layer, an intermediate layer and an output layer connected in sequence; in the intermediate layer of the fully connected feedforward network FFN, the number of neurons is reduced by pruning. For the fully connected feedforward network FFN of the layer Transformer Encoder, there are features Z l , the characteristic of each neuron j is The intermediate layer is implemented by a hidden layer; By calculating the Gram matrix and dependency score Evaluate the importance score of each neuron; Gram Matrix The elements in are defined as Among them, k is a kernel function used to calculate the similarity between two vectors; Dependency score The formula is: in, is defined in The Gram matrix on , tr represents the trace of the matrix, CLC is a constant matrix; the dependency score is the importance score of each neuron; The neuron pruning ratio k obtained by searching with genetic algorithm l , retain the ones with the highest importance scores neurons; the rest were deleted; Pruning method of Multi-head Attention layer: In the multi-head attention mechanism layer Multi-head Attention of the l-th layer Transformer Encoder, the output feature of each attention head h is represented as Calculate the Gram matrix K of each attention head l and dependency score Evaluate the importance score of each head; According to the pruning ratio σ obtained by the search l , the top σ with the highest score l ·N h Attention heads are retained and the rest are deleted. h is the total number of attention heads.

8. Image classification system based on pruned deep learning model, characterized by: include: A pruning module is configured to: obtain a deep learning model, perform a pruning operation 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: using a genetic algorithm to perform a 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 also includes: 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 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, wherein the training set is an image with a known image classification label; input the training set into the deep learning model, train the model, and obtain a trained deep learning model; copy the model parameters of the trained deep learning model into the pruned deep learning model, and obtain a trained pruned deep learning model; The classification module is configured to: obtain the image to be classified, input the image to be classified into the trained pruned deep learning model, and obtain the classification result of the image.

9. An electronic device, comprising: a memory for non-transitory storage of computer readable instructions; as well as a processor for executing the computer readable instructions, When the computer-readable instructions are executed by the processor, the method described in any one of claims 1 to 7 is executed.

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

Images