Hardware-friendly pruning method based on Vision Transform neural network

Abstract

The invention discloses a hardware-friendly pruning method based on a Vision Transform neural network, and belongs to the technical field of calculation, reckoning or counting. The pruning method is a scheme combining layer-by-layer N: M weight pruning and non-parameterized multi-stage adaptive token pruning, and comprises the following steps: calculating an average Hessian trace of each layer of the network according to a neural network and a pre-training weight, realizing layer-by-layer N: M weight pruning according to the distribution of the average Hessian trace, carrying out weight pruning retraining and updating the weight, and carrying out multi-stage adaptive token pruning. And then deploying an adaptive token pruning module, carrying out token pruning retraining and weight updating, and finally deploying the neural network adopting the pruning method to a hardware platform. According to the method, the parameter quantity and the calculation quantity of the ViT network are effectively reduced by combining hardware-friendly weights and tokens with a pruning method, the reasoning delay and the hardware overhead of the ViT are reduced, high precision is kept, and an optimization method is provided for deployment of the ViT network at a resource-limited edge end.

Description

Technical Field

[0001] The present invention relates to neural network compression and acceleration technology, and specifically discloses a hardware-friendly pruning method based on a Vision Transformer neural network, belonging to the technical field of calculation, inference or counting. Background Art

[0002] In recent years, the development of deep learning has promoted the progress of various computer vision tasks such as image classification and object detection. Among the current neural networks, the Vision Transformer (ViT) architecture has demonstrated excellent performance and flexibility. Compared with traditional convolutional neural networks, ViT can better capture the global information in the image through the self-attention mechanism, thus achieving state-of-the-art performance in many visual tasks.

[0003] However, the high performance of ViT neural network is accompanied by high computational complexity and high storage requirements: there are a large number of matrix multiplications in the network, and the computational complexity of the multi-head self-attention mechanism is quadratic relative to the number of input tokens. The number of network parameters is also much higher than that of traditional convolutional neural networks. These characteristics of ViT neural network limit its application in resource-constrained embedded systems, mobile devices and other edge devices.

[0004] Pruning is an effective solution to the above problems. Removing redundant weights in the network can significantly reduce the number of parameters and computations. Removing redundant tokens in the image can quadratically reduce the computation of the multi-head self-attention module, thereby improving the model's reasoning speed and reducing hardware overhead. However, existing research mainly includes the following three types of token pruning schemes: the first type of token pruning scheme directly prunes the multi-head self-attention module, which is an extremely coarse-grained pruning that seriously reduces the network accuracy; the second type of token pruning scheme introduces neural networks and nonlinear functions to achieve token recognition and pruning, but this will bring additional parameters, reasoning delays and hardware overhead, which is not friendly to resource-constrained edge devices; the third type of token pruning scheme performs fixed pruning of tokens, and the number of tokens pruned for images of different complexities is the same. This pruning method will prune the tokens corresponding to objects in complex images, which has a greater impact on network accuracy. In the current weight pruning method, all layers of the network use the same pruning rate, and the network accuracy loss is large. In addition, in order to reduce the network computation and parameter amount while maintaining good network accuracy, which is an urgent problem to be solved, there are currently few studies that consider token pruning and weight pruning at the same time.

[0005] Therefore, the present invention aims to propose a hardware-friendly pruning method based on Vision Transformer neural network to overcome the above defects. Summary of the invention

[0006] The purpose of the present invention is to address the deficiencies of the above-mentioned background technology and provide a hardware-friendly pruning method based on VisionTransformer neural network, which combines layer-by-layer N:M weight pruning and adaptive token pruning to solve the technical problems of high computational complexity and high storage requirements of ViT neural network, and achieve the purpose of reducing the inference delay of ViT neural network, reducing the hardware overhead of edge deployment, and improving the inference speed and energy efficiency of the network in a resource-constrained environment.

[0007] The present invention adopts the following technical solutions to achieve the above-mentioned invention object:

[0008] A hardware-friendly pruning method based on Vision Transformer neural network includes the following steps:

[0009] Step 1, calculate the average Hessian trace of each layer of the neural network;

[0010] Step 2: Perform N:M weight pruning layer by layer according to the distribution of the average Hessian trace of each layer, where N and M are positive integers and N <M;

[0011] Step 3, retrain the network after weight pruning in step 2 and update the weights;

[0012] Step 4, performing non-parametric multi-stage adaptive token pruning on the network retrained in step 3;

[0013] Step 5: Retrain the network after token pruning in step 4 to obtain the Vision Transformer neural network, and deploy the Vision Transformer neural network to the hardware platform.

[0014] As a further optimization scheme of the hardware-friendly pruning method based on the Vision Transformer neural network, step 1 calculates the average Hessian trace of each layer of the neural network according to the pre-trained weights. The average Hessian trace of each layer includes: the average Hessian trace of the multi-head self-attention module and the average Hessian trace of the feedforward module. The average Hessian trace of the multi-head self-attention module includes the average Hessian trace of the input mapping and the average Hessian trace of the output mapping. The average Hessian trace of the feedforward module includes the average Hessian trace of two fully connected layers.

[0015] As a further optimization scheme of the hardware-friendly pruning method based on the Vision Transformer neural network, the specific method of performing N:M weight pruning layer by layer according to the distribution of the average Hessian traces of each layer in step 2 is as follows: according to the distribution of the average Hessian traces of the multi-head self-attention modules in each layer, determine the optimal numerical interval division method of the average Hessian traces of the multi-head self-attention modules in each layer, and configure N:M values ​​for each numerical interval of the average Hessian traces of the multi-head self-attention modules in each layer according to the optimal numerical interval division method; use the same method to configure N:M values ​​for each numerical interval of the average Hessian traces of the feedforward modules in each layer; and prune the weight matrices of the multi-head self-attention modules in each layer and the weight matrices of the feedforward modules in each layer by column according to the configured N:M.

[0016] As a further optimization scheme of the hardware-friendly pruning method based on the Vision Transformer neural network, the specific method of step 4 for performing non-parametric multi-stage adaptive token pruning on the network retrained in step 3 is as follows: a token pruning layer is selected during the network forward inference process, and each image token score and the pruning threshold of the token pruning layer are calculated according to the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer, and the image token scores of the token pruning layer are sorted in descending order, and partial sum accumulation operations are performed in sequence on the descending sorted image token scores, and the accumulation operation is stopped when the partial sum obtained by accumulation is greater than the pruning threshold of the token pruning layer, and a pruning decision instruction is generated according to the image token that does not participate in the partial sum accumulation operation, and the pruning decision instruction is executed on the image token of the layer normalization module input to the feedforward module of the token pruning layer.

[0017] As a further optimization scheme of the hardware-friendly pruning method based on the Vision Transformer neural network, the weight matrices of the multi-head self-attention modules of each layer and the weight matrices of the feedforward modules of each layer are pruned column by column according to the configured N:M. Specifically, the weights of each column are grouped according to the value of M, and the largest N weights in each group of weights are kept as they are and the other weights are assigned to zero.

[0018] As a further optimization scheme of the hardware-friendly pruning method based on the Vision Transformer neural network, the optimal numerical interval division method of the average Hessian trace of each layer of the multi-head self-attention module and the optimal numerical interval division method of the average Hessian trace of each layer of the feedforward module are determined according to the network accuracy and weight pruning benefit.

[0019] As a further optimization scheme of a hardware-friendly pruning method based on the Vision Transformer neural network, the token score of each image and the pruning threshold of the token pruning layer are calculated according to the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer. Specifically, the first row of data of the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer is summed to obtain a new row of data, and the remaining elements of the new row of data except the first element correspond to the token scores of each image in the token pruning layer. The token scores of each image in the token pruning layer are summed and multiplied by the scaling factor to obtain the pruning threshold of the token pruning layer.

[0020] As a further optimization of the hardware-friendly pruning method based on the Vision Transformer neural network, the pruning threshold of the token pruning layer is Among them, Threshold is the pruning threshold of the token pruning layer, α is the pruning ratio factor, 0<α<1, token score i is the score of the i-th image token, and N is the number of image tokens input to the token pruning layer.

[0021] As a further optimization scheme of a hardware-friendly pruning method based on the Vision Transformer neural network, the Vision Transformer neural network obtained in step 5 includes at least one layer of network structure, each layer of the network structure includes a multi-head self-attention module and a feedforward module that executes computing tasks in serial, the token pruning layer also includes an adaptive token pruning module that executes computing tasks in parallel with the multi-head self-attention module of this layer, the adaptive token pruning module receives the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer, and the adaptive token pruning module outputs a pruning decision instruction to the input end of the layer normalization module before the feedforward module in the token pruning layer.

[0022] As a further optimization scheme of the hardware-friendly pruning method based on the Vision Transformer neural network, the adaptive token pruning module includes: a token score and pruning threshold calculation unit, a token score descending sorting unit, and a token score part and calculation unit;

[0023] A token score and pruning threshold calculation unit, whose input end receives the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer, and calculates the token score of each image and the pruning threshold of the token pruning layer according to the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer;

[0024] A token score descending sorting unit, whose input end receives each image token score of the token pruning layer and sorts each image token score of the token pruning layer in descending order;

[0025] A token score partial and calculation unit, whose input end receives the pruning threshold of the token pruning layer and the scores of each image token after descending sorting, performs partial and accumulation operations on each image token score after descending sorting in sequence, stops the accumulation operation when the partial sum obtained by accumulation is greater than the pruning threshold of the token pruning layer, generates pruning decision instructions according to the image tokens that do not participate in the partial and accumulation operation, and outputs the pruning decision instructions to the input end of the layer normalization module before the feedforward module in the token pruning layer.

[0026] The present invention adopts the above technical solution and has the following beneficial effects:

[0027] (1) The present invention provides a hardware-friendly pruning method based on the Vision Transformer neural network, which combines weight pruning and token pruning, can simultaneously reduce the computational complexity and parameter quantity of the ViT network, reduce the network's inference delay and hardware overhead while maintaining high accuracy, and provide an optimization method for the deployment of the network at the edge end with limited resources.

[0028] (2) Considering that the weights of different layers of the ViT network have different sensitivities to errors, the present invention adopts a layer-by-layer N:M weight pruning method to perform weight pruning, applies the Hessian trace to weight pruning, and has different pruning rates for weights of different layers and different modules in the same layer, which effectively reduces the amount of network parameters and storage overhead while reducing the accuracy loss caused by pruning. In addition, the adopted N:M pruning method is very friendly to hardware acceleration.

[0029] (3) The token pruning method adopted in the present invention is an adaptive token pruning scheme. The number of tokens pruned for images of different complexities is different, which is more in line with real-life images and application scenarios. In addition, the token pruning method adopted uses the intermediate value of ViT network inference to achieve token pruning, without introducing additional networks, parameters and nonlinear functions, and is very friendly to hardware acceleration.

[0030] (4) The implementation process of the pruning method proposed in the present invention is simple, the functions of each step are clearly distinguished, and it can be easily applied to various ViT networks and related variants. BRIEF DESCRIPTION OF THE DRAWINGS

[0031] Figure 1 This is an overall flow chart of the pruning method based on the ViT neural network adopted in the present invention.

[0032] Figure 2 The specific structure of a layer of encoder in the ViT neural network that does not contain the adaptive token pruning module involved in the present invention.

[0033] Figure 3 The present invention discloses a specific structure of an encoder layer in a ViT neural network including the adaptive token pruning module involved in the present invention.

[0034] Figure 4 It is the average Hessian trace of the multi-head self-attention module in all layers calculated according to the network structure and pre-trained weights in the example of the present invention.

[0035] Figure 5 This is a visualized pruning result of multi-stage adaptive token pruning in an example of the present invention. DETAILED DESCRIPTION

[0036] The technical solution of the invention is described in detail below in conjunction with the accompanying drawings. It should be understood that these embodiments are only used to illustrate the invention and are not used to limit the scope of the invention. After reading the invention, various equivalent modifications made to the invention by those skilled in the art all fall within the scope defined by the claims attached to the invention.

[0037] The hardware-friendly pruning method based on the Vision Transformer neural network involved in the present invention uses a layer-by-layer N:M weight pruning method to remove redundant weights in the network, and uses a non-parametric multi-stage adaptive token pruning method to dynamically remove redundant tokens in the image during reasoning, and combines the two pruning methods to achieve hardware-friendly weight-token pruning. In the process of weight-token pruning, different weight pruning parameters and token pruning parameters are selected for each layer according to the accuracy of the network and the pruning benefit, and the model is retrained to restore the accuracy after the pruning parameters are determined. Specifically, it includes the following 5 steps.

[0038] Step 1: Calculate the average Hessian trace of each layer of the neural network

[0039] The average Hessian trace of each layer includes the average Hessian trace of the multi-head self-attention module and the average Hessian trace of the feedforward module. The average Hessian trace of the multi-head self-attention module includes the average Hessian trace of the input mapping and the average Hessian trace of the output mapping. The average Hessian trace of the feedforward module includes the average Hessian trace of two fully connected layers. The structure of a layer encoder of the ViT network is as follows Figure 2 As shown, the multi-head self-attention module has two linear mappings, input mapping and output mapping. The feed-forward module has two fully connected layers. The pre-trained weights of the network are used to calculate the average Hessian trace of the input mapping in each encoder layer, the average Hessian trace of the output mapping, and the average Hessian trace of each of the two fully connected layers.

[0040] Step 2: Perform layer-by-layer N:M weight pruning based on the distribution of the average Hessian trace of each layer

[0041] N:M weight pruning means that the largest N of every M weights remain the original value, and the remaining weights are pruned to zero. The value of N:M represents the weight pruning rate and each layer can take a different value of N:M.

[0042] The Hessian trace is very effective in reflecting the quantization error of weights, and both weight quantization and weight pruning are related to error sensitivity. Therefore, the application of the Hessian trace is extended to weight pruning, and the value of N:M of each layer is determined by the size of the average Hessian trace of the layer. For a network with several layers, the higher the average Hessian trace of the layer, the more sensitive the weight is to error, and the higher the N:M selected, that is, the lower the pruning rate, and vice versa.

[0043] When performing N:M weight pruning on the multi-head self-attention module layer by layer, the average Hessen trace of the multi-head self-attention module is divided into several numerical intervals according to the calculated distribution of the average Hessen trace of the multi-head self-attention module. Each group of numerical intervals corresponds to an N:M value. The higher the average Hessen trace value corresponding to the numerical interval, the smaller the weight pruning rate; then the N:M value is selected for each layer of the feedforward module in the same way as the multi-head self-attention module. According to the value of N:M, the weight matrix of the corresponding module is pruned by column, that is, the weight matrix of each layer of the multi-head self-attention module is pruned by column according to the N:M configured according to the average Hessen trace of the multi-head self-attention module, and the weight matrix of each layer of the feedforward module is pruned by column according to the N:M configured according to the average Hessen trace of the feedforward module: the largest N of each M weights are kept original, and the other values ​​are set to zero; then the pruned weights are used to count the network accuracy and the amount of weight pruning, and the optimal average Hessen trace numerical interval division method and N:M value are determined according to the accuracy and weight pruning benefits.

[0044] After obtaining the above-mentioned optimal average Hessen trace value interval division method and the N:M value of each module in each layer, the weight matrix of each module in each layer is pruned column by column according to the determined N:M value.

[0045] Step 3: Retrain the model to restore accuracy and update weights, keeping the pruned weights at zero during training.

[0046] Step 4: Token pruning using a non-parametric multi-stage adaptive token pruning method

[0047] Based on the weights retrained by layer-by-layer N:M weight pruning, non-parametric multi-stage adaptive token pruning is performed. Token pruning is performed in multiple stages during the network forward reasoning process and uses the intermediate values ​​of network reasoning without introducing additional networks and parameters. That is, token pruning is performed on several layers of the network during the reasoning process, and token pruning is not performed on other layers. Adaptive pruning with different numbers of pruning tokens is performed for images of different complexity.

[0048] In layers where token pruning is not required, directly Figure 2 The forward inference process of the network is performed in the encoder shown in . In the layer where token pruning is required, the network is as follows Figure 3 Forward inference is performed in the ViT encoder with an adaptive token pruning module shown in the figure. The adaptive token pruning module and the multi-head self-attention module of this layer are calculated in parallel. The adaptive token pruning module includes: a token score and pruning threshold calculation unit, a token score descending sorting unit, a token score part and a calculation unit. First, the output matrix of the Softmax function of all heads of the multi-head self-attention module of this layer is input into the adaptive token pruning module. The token score and pruning threshold calculation unit extracts the first row of data of the Softmax function output matrix of all heads, and sums the first row of data of the Softmax function output matrix of all heads to obtain a new row of data. After removing the first element of the new row of data, the remaining elements correspond to the scores of each image token in this layer, as shown in the following formula. The scores of all image tokens of this layer are summed and multiplied by a scaling factor α (0<α<1) to obtain the pruning threshold:

[0049]

[0050] Among them, Threshold is the pruning threshold, α is the pruning ratio factor, and token score i is the score of the i-th image token, and N is the number of image tokens input in the pruned layer. Then, the token score descending sorting unit sorts the scores of the image tokens in descending order, and the token score partial sum calculation unit adds the sorted token scores from high to low one by one and obtains the partial sum of the token scores. After each new token score is added to the partial sum of the token scores, the partial sum is compared with the pruning threshold. If the partial sum of the token scores is less than the threshold, the next token score is added to the partial sum and the above process is repeated. If the partial sum of the token scores is greater than the threshold, the addition is stopped, and the image tokens corresponding to the token scores that have been added to the partial sum of the token scores are retained, while the image tokens corresponding to the token scores that have not been added are pruned. After the token score partial sum calculation unit completes the partial sum solution and the threshold comparison, a pruning decision instruction is generated according to the image tokens corresponding to the token scores that have not been added, and the image tokens input to the layer normalization module before the feedforward module in the layer execute the pruning decision instruction. The image tokens corresponding to the token scores that have not been added no longer participate in the layer normalization before the feedforward module of the layer, the calculation of the feedforward module of the layer and the remaining layers of the network.

[0051] In the layer that needs token pruning, the adaptive token pruning module starts to calculate after the network calculates the Softmax function. The above pruning process is the same as Figure 3The second matrix multiplication of the multi-head self-attention module and the output mapping are calculated in parallel, without introducing additional inference delay. For the scores of the above image tokens, the distribution of token scores for images of different complexity is different. For images with smaller objects to be classified, the scores of tokens are concentrated in a few tokens, and the scores of different tokens vary greatly. Fewer image tokens can exceed the pruning threshold, and the number of pruned tokens is high; for images with larger objects to be classified, the scores of tokens are more dispersed, and the scores of different tokens are smaller and the differences are smaller. More image tokens can exceed the pruning threshold, and the number of pruned tokens is low. At this time, the number of pruned tokens in different images is different, realizing adaptive pruning.

[0052] Step 5: Retrain the network after token pruning and update the weights

[0053] Retrain the model to restore accuracy and update weights while keeping the zero-valued weights fixed in N:M weight pruning, obtain the Vision Transformer neural network, and deploy the Vision Transformer neural network to the hardware platform.

[0054] The layer-by-layer N:M weight pruning method has the following characteristics: the weight pruning method is a static structural pruning scheme, the weights have been pruned before the network is inferred, which is very friendly to hardware deployment. It also takes into account the different sensitivity of each layer of the network to errors, and the weight pruning rate used in each layer is different, which can minimize the accuracy loss of the network.

[0055] The multi-stage token adaptive pruning method has the following characteristics: the token pruning method is a dynamic structural pruning scheme, the token pruning is related to the input image and is completed during the network reasoning process, the pruning uses the intermediate values ​​during network reasoning and the pruning process is parallel to the network reasoning process, no additional parameters and reasoning delays are introduced, and the network reasoning can be effectively accelerated during hardware deployment, and the adaptive scheme can achieve different token pruning rates according to the complexity of the image, which is close to the application scenarios in reality.

[0056] The pruning method of the present invention is described in detail by taking an example as follows.

[0057] We take a variant of the ViT neural network, the DeiT-S network, as an example. The layer structure of the network is exactly the same as that of ViT, but the difference is that the knowledge distillation training strategy is adopted. The DeiT-S neural network has 12 layers, and each layer's multi-head self-attention module consists of 6 heads. The number of image tokens is 196 when not pruned.

[0058] First, build the DeiT-S neural network and calculate the average Hessian traces of all layers of the network. First, build the original DeiT-S neural network without pruning, and use the DeiT-S neural network and its pre-trained weights to calculate the average Hessian traces of all layers in the network. The average Hessian trace of each layer includes the average Hessian trace of the multi-head self-attention module and the average Hessian trace of the feedforward module. The average Hessian trace of the multi-head self-attention includes the average Hessian trace of the input mapping and the average Hessian trace of the output mapping. The average Hessian trace of the feedforward module includes the average Hessian traces of the two fully connected layers.

[0059] Determine the value of N:M when pruning the weights of each layer. Count the average Hessian traces of all layers of the DeiT-S neural network and get the following Figure 4 The average Hessian trace distribution of the multi-head self-attention module is shown in Figure 2. In this example, the two groups of average Hessian traces of the multi-head self-attention module are divided into three value intervals, namely [0, Th 1 ]、[Th 1 ,Th 2 ]、[Th 2 ,0.001], the N:M values ​​of the three numerical intervals are 1:8, 1:4, and 1:2 respectively. Similarly, in this example, the two groups of average Hessen traces of the feedforward module are divided into three numerical intervals, namely [0,Th 3 ]、[Th 3 ,Th 4 ]、[Th 4 ,0.001], the N:M values ​​of the three numerical intervals are 1:8, 1:4, and 1:2 as examples. Different numerical interval combinations are selected for the average Hessian trace of the DeiT-S neural network, and the value of N:M is determined for each module in each layer. The accuracy and weight pruning amount of the network are statistically analyzed, and the results are shown in Table 1:

[0060] <![CDATA[Th 1 ]]> <![CDATA[Th 2 ]]> <![CDATA[Th 3 ]]> <![CDATA[Th 4 ]]> Weight pruning amount Accuracy 0.1 0.2 0.1 0.2 68.9％ 79.5％ 0.1 0.3 0.1 0.3 71.7％ 79.7％ 0.2 0.4 0.1 0.3 74.4％ 79.0％ 0.2 0.4 0.2 0.4 76.1％ 78.8％ 0.1 0.3 0.2 0.4 73.4％ 79.3％ 0.3 0.5 0.3 0.5 79.3％ 78.2％

[0061] Table 1

[0062] According to the network accuracy and weight pruning amount in Table 1, the parameter values ​​in this example are: Th 1 =0.1,Th 2 =0.3,Th 3 =0.1,Th 4 =0.3, and determine the N:M values ​​corresponding to the four weight groups of each layer according to the numerical range. It should be noted that the above N:M values ​​and combinations are not limited to 1:8, 1:4 and 1:2, and multiple values ​​such as 2:4, 3:4 and multiple N:M combinations can also be selected.

[0063] Perform weight pruning and retraining. Perform column-wise weight pruning according to the determined N:M: every M weights in each column are grouped together, the N largest weights among these M weights are kept as they are, and the other weights are pruned and assigned to zero. Retrain the model to restore accuracy and update the weights, and keep the pruned weights to zero.

[0064] The adaptive token pruning module is deployed in the DeiT-S neural network to realize the adaptive pruning of image tokens. The score and pruning threshold of the image token are calculated. In the DeiT-S neural network of this example, the token pruning positions are selected as the 4th, 7th, and 10th layers. First, it is determined whether the current network reasoning is in the token pruning layer. If not, no token pruning is performed. If it is, the token adaptive pruning process is performed: the output matrix of the Softmax function of all heads in the multi-head self-attention module of this layer is sent to the token adaptive pruning module. In the token score and pruning threshold calculation unit, the first row of data of the output matrix of the Softmax function of all heads is extracted and added to obtain a new row of data. After removing the first element in the new data of this row, the remaining elements correspond to the scores of all image tokens in this layer. The scores of all image tokens in this layer are summed and multiplied by a scaling factor α (0<α<1) to obtain the pruning threshold. In this example, the scaling factor α can be 0.95, 0.93, 0.85, 0.8, etc. Different scaling factors α can obtain different network accuracy and calculation reduction benefits.

[0065] The token scores are sorted in descending order and the partial sum of the token scores is calculated. In the token pruning layer, the scores of the above image tokens are sorted in descending order in the token score descending sorting unit, and then the sorted token scores are added one by one from high to low in the token score partial sum calculation unit to obtain the partial sum of the token scores. After each new token score is added to the partial sum of the token scores, the partial sum is compared with the above pruning threshold. If the partial sum of the token scores is less than the threshold, the next token score is added to the partial sum and the above process is repeated. If the partial sum of the token scores is greater than the threshold, the addition is stopped, and the image tokens corresponding to the token scores that have been added to the partial sum are retained, while the image tokens whose token scores have not been added are directly discarded after the calculation of the multi-head self-attention module is completed, and no longer participate in the calculation of the feedforward module of this layer and the remaining layers of the network. When α=0.9 is used and token pruning is performed on the 4th, 7th, and 10th layers, the visualization results of multi-stage adaptive token pruning are shown as follows: Figure 5 As shown, the shaded part represents the tokens that are pruned at this layer.

[0066] Retraining with adaptive token pruning is performed to restore accuracy and update weights. An adaptive token pruning module is added to the DeiT-S neural network. The model is retrained again to restore accuracy and update weights while maintaining the retrained weights of N:M weight pruning. During the training process, the pruned weights are kept at zero.

[0067] The DeiT-S neural network using the pruning method proposed in this invention is deployed on a hardware platform to evaluate the accuracy of the network and the benefits of pruning. The above-mentioned layer-by-layer N:M weight pruning and non-parametric multi-stage adaptive token pruning are both hardware-friendly pruning schemes that can be easily deployed in hardware. 1 =0.1,Th 2 =0.3,Th 3 =0.1,Th 4 When α=0.3 and α=0.9, the parameters of the DeiT-S neural network in this example are evaluated to be reduced by 71.7%, and the network inference speed is increased by 1.6 times.

[0068] The above shows and describes the basic principles, main features and advantages of the present invention. Those skilled in the art should understand that the present invention is not limited to the above specific embodiments. The above specific embodiments and the description in the specification are only for further illustrating the principles and preparation effects of the present invention. Without departing from the spirit and scope of the present invention, the present invention may have various changes and improvements, which fall within the scope of the present invention. The scope of the present invention is defined by the claims and their equivalents.

Claims

1. A hardware-friendly pruning method based on Vision Transformer neural network, characterized in that: The steps include: Step 1, calculate the average Hessian trace of each layer of the neural network; Step 2: Perform N:M weight pruning layer by layer according to the distribution of the average Hessian trace of each layer, where N and M are positive integers and N <M; Step 3, retrain the network after weight pruning in step 2 and update the weights; Step 4, performing non-parametric multi-stage adaptive token pruning on the network retrained in step 3; Step 5: retrain the network after token pruning in step 4 to obtain the Vision Transformer neural network, and deploy the Vision Transformer neural network to the hardware platform.

2. According to claim 1, a hardware-friendly pruning method based on Vision Transformer neural network is characterized in that: The step 1 calculates the average Hessian trace of each layer of the neural network according to the pre-trained weights, and the average Hessian trace of each layer includes: the average Hessian trace of the multi-head self-attention module and the average Hessian trace of the feedforward module, the average Hessian trace of the multi-head self-attention module includes the average Hessian trace of the input mapping and the average Hessian trace of the output mapping, and the average Hessian trace of the feedforward module includes the average Hessian trace of two fully connected layers.

3. According to claim 2, a hardware-friendly pruning method based on Vision Transformer neural network is characterized in that: The specific method of step 2 for performing layer-by-layer N:M weight pruning according to the distribution of the average Hessian traces of each layer is as follows: according to the distribution of the average Hessian traces of the multi-head self-attention modules of each layer, determine the optimal numerical interval division method of the average Hessian traces of the multi-head self-attention modules of each layer, and configure N:M values ​​for each numerical interval of the average Hessian traces of the multi-head self-attention modules of each layer according to the optimal numerical interval division method; use the same method to configure N:M values ​​for each numerical interval of the average Hessian traces of the feedforward modules of each layer; and prune the weight matrices of the multi-head self-attention modules of each layer and the weight matrices of the feedforward modules of each layer by column according to the configured N:M.

4. According to claim 3, a hardware-friendly pruning method based on Vision Transformer neural network is characterized in that: The specific method of step 4 for performing non-parametric multi-stage adaptive token pruning on the network retrained in step 3 is as follows: selecting a token pruning layer during the network forward inference process, calculating each image token score and the pruning threshold of the token pruning layer according to the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer, sorting the image token scores of the token pruning layer in descending order, performing partial sum accumulation operations on the descending sorted image token scores in sequence, stopping the accumulation operation when the partial sum obtained by accumulation is greater than the pruning threshold of the token pruning layer, generating pruning decision instructions according to the image tokens that do not participate in the partial sum accumulation operation, and executing the pruning decision instructions on the image tokens of the layer normalization module input to the feedforward module of the token pruning layer.

5. According to claim 3, a hardware-friendly pruning method based on Vision Transformer neural network is characterized in that: According to the configured N:M, the weight matrices of the multi-head self-attention modules of each layer and the weight matrices of the feedforward modules of each layer are pruned by columns. Specifically, the weights of each column are grouped according to the value of M, and the largest N weights in each group of weights are kept as they are and the other weights are assigned to zero.

6. According to claim 3, a hardware-friendly pruning method based on Vision Transformer neural network is characterized in that: According to the network accuracy and weight pruning benefits, the optimal numerical interval division method of the average Hessian trace of each layer of the multi-head self-attention module and the optimal numerical interval division method of the average Hessian trace of each layer of the feedforward module are determined.

7. A hardware-friendly pruning method based on Vision Transformer neural network according to claim 4, characterized in that: The token score of each image and the pruning threshold of the token pruning layer are calculated according to the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer. Specifically, the first row of data of the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer is summed to obtain a new row of data, wherein the remaining elements of the new row of data except the first element correspond to the token scores of each image in the token pruning layer, and the sum of the token scores of each image in the token pruning layer is multiplied by the proportional factor to obtain the pruning threshold of the token pruning layer.

8. A hardware-friendly pruning method based on Vision Transformer neural network according to claim 7, characterized in that: The pruning threshold of the token pruning layer is Among them, Threshold is the pruning threshold of the token pruning layer, α is the pruning ratio factor, 0<α<1, token score i is the score of the i-th image token, and N is the number of image tokens input to the token pruning layer.

9. A hardware-friendly pruning method based on Vision Transformer neural network according to any one of claims 1 to 8, characterized in that: The Vision Transformer neural network obtained in step 5 includes at least one layer of network structure, and each layer of the network structure includes a multi-head self-attention module and a feedforward module that execute computing tasks in series. The token pruning layer also includes an adaptive token pruning module that executes computing tasks in parallel with the multi-head self-attention module of this layer. The adaptive token pruning module receives the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer, and the adaptive token pruning module outputs a pruning decision instruction to the input end of the layer normalization module before the feedforward module in the token pruning layer.

10. A hardware-friendly pruning method based on Vision Transformer neural network according to claim 9, characterized in that: The adaptive token pruning module comprises: A token score and pruning threshold calculation unit, whose input end receives the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer, and calculates the token score of each image and the pruning threshold of the token pruning layer according to the Softmax function output matrix of all heads of the multi-head self-attention module in the token pruning layer; a token score descending sorting unit, whose input end receives each image token score of the token pruning layer and sorts each image token score of the token pruning layer in descending order; and, A token score partial and calculation unit, whose input end receives the pruning threshold of the token pruning layer and the scores of each image token after descending sorting, performs partial and cumulative operations on each image token score after descending sorting in sequence, stops the accumulation operation when the partial sum obtained by accumulation is greater than the pruning threshold of the token pruning layer, generates pruning decision instructions according to the image tokens that do not participate in the partial and cumulative operation, and outputs the pruning decision instructions to the input end of the layer normalization module before the feedforward module in the token pruning layer.

Images