A model pruning method and system based on weighted attention alignment

Abstract

The application discloses a model pruning method and system based on weighted attention alignment, relates to artificial intelligence, and is proposed in view of the lack of correlation in the prior art. The method comprises the following steps: S1. pre-tuning: first, the target model is initialized by using the pre-training model parameters, and then the L1 regularization fine-tuning is performed on the batch normalization layer parameters by using the data of the target domain for several times; S2. pre-pruning: the absolute values of the batch normalization layer parameters of the trained target model are globally sorted, the tail channels are pruned, and a pre-pruned model is obtained; S3. model pruning training based on weighted attention alignment: the pruned model is fine-tuned and trained, the pre-training model is introduced by using the weighted attention alignment mechanism, and a pruning model for the target domain data set is obtained. The method has the advantages that the pre-training knowledge can be efficiently utilized, the pruning and migration can be effectively combined, the pruning performance can be improved, and the training time can be reduced.

Description

Technical Field

[0001] The present invention relates to artificial intelligence technology, and in particular to a model pruning method and system based on weighted attention alignment. Background Art

[0002] In recent years, deep learning has rapidly developed, impacting various fields including computer vision, natural language processing, and speech recognition. Convolutional neural networks (CNNs) for images have demonstrated powerful feature extraction capabilities, becoming a foundational model for image recognition tasks. To improve model performance, CNNs have become deeper, wider, and more complex, leading to a dramatic increase in model storage space and computational complexity. In response, a growing number of researchers are studying model compression, specifically how to use various techniques to reduce the size, complexity, and computational complexity of deep models, improve the speed and efficiency of inference, and maintain performance as much as possible for deployment and operation on resource-constrained devices. Model pruning can accelerate neural networks by removing unimportant components, while retaining a high degree of accuracy and being simple and effective. Filter pruning has garnered significant attention due to its applicability to any convolutional model architecture and its deployment-friendly nature.

[0003] However, the pruned model's generalization performance is significantly weakened. Specifically, while the performance of a large model pre-trained on ImageNet does not significantly decrease after pruning on ImageNet, when migrating it to a smaller dataset, the performance drops significantly compared to migrating the full model. This means that the pruned model is more sensitive to distribution differences between datasets and cannot be effectively migrated to other datasets. In other words, only pre-training the full model can improve the performance of the pruned model on a smaller dataset.

[0004] To achieve this goal, some existing methods, such as Liu et al. (Liu B, Cai Y, Guo Y, et al. TransTailor: Pruning the Pre-trained Model for Improved Transfer Learning. [Z]. 20218627-8634), propose a fine-tuning-then-pruning algorithmic process: the pre-trained model is first fine-tuned on a small dataset, and then the fine-tuned model is pruned during training. Pruning is mainly based on Taylor expansion, cumulative activation sums, and attention mechanisms. However, existing methods have the following problems: First, while pre-trained knowledge is fundamental to improving performance, these methods simply use pre-trained knowledge as a simple initialization and fail to fully utilize the abundant useful knowledge in the pre-trained model. Second, pruning and transfer are not effectively combined. Pruning naturally ranks feature channels by importance, while transfer in existing methods only involves fine-tuning and does not involve feature alignment. Third, these methods simply perform transfer and pruning training sequentially or even iteratively, which consumes a lot of computational and time costs. Summary of the Invention

[0005] The purpose of the present invention is to provide a model pruning method and system based on weighted attention alignment to solve the problems existing in the above-mentioned prior art.

[0006] The model pruning method based on weighted attention alignment described in the present invention includes the following steps:

[0007] S1. Pre-fine-tuning: Initialize the target model using the pre-trained model parameters, and then fine-tune the batch normalization layer parameters several times with L1 regularization using data from the target domain.

[0008] S2. Pre-pruning: Globally sort the absolute values ​​of the batch normalization layer parameters of the trained target model and prune the last channels to obtain a pre-pruned model.

[0009] S3. Model pruning training based on weighted attention alignment: Prune and fine-tune the pruned model, introduce guidance from the pre-trained model through the weighted attention alignment mechanism, and obtain a pruned model for the target domain dataset.

[0010] The present invention provides a model pruning system based on weighted attention alignment, which uses the model pruning method to prune the model to obtain a pruned model for a small data set.

[0011] The model pruning method and system based on weighted attention alignment described in the present invention have the advantages of overcoming the shortcomings of the existing technology:

[0012] I. Weighted attention alignment: To address the problem of insufficient utilization of pre-training knowledge, first calculate the weighted attention feature map through the parameters of the batch normalization layer, and then align the weighted attention maps of the large model and the small model to realize efficient utilization of pre-training knowledge.

[0013] II. Combination of pruning and migration: To address the problem that pruning and migration are not related, the batch normalization layer parameters are used to effectively combine the two. The batch normalization layer parameters are a quantitative evaluation of the importance of each channel, which is not only used as a basis for channel pruning, but also used for the calculation of the weighted attention feature map, realizing the effective combination of pruning and migration.

[0014] III. Optimization of joint training strategy: To address the problem of step-by-step and iterative algorithm training, the algorithm flow is designed to migrate and prune jointly, and the migration of weighted attention is performed at the same time as the pruning fine-tuning training, which not only improves the pruning performance, but also significantly reduces the time cost of the training process. BRIEF DESCRIPTION OF DRAWINGS

[0015] Figure 1 is a schematic diagram of the model pruning method described in the present application.

[0016] Figure 2 is a schematic diagram of the calculation process of the weighted attention map described in the present application. DETAILED DESCRIPTION

[0017] The present application can fully utilize pre-training knowledge and improve the performance of the pruned model for small data sets by calculating the weighted attention map of the batch normalization layer parameters and applying alignment constraints. At the same time, it effectively combines pruning and migration through the batch normalization layer parameters, improves the interpretability of the method, and performs migration of weighted attention at the same time as pruning fine-tuning training, which not only improves the pruning performance, but also significantly reduces the time cost of the existing step-by-step or even iterative training process.

[0018] A model pruning method based on weighted attention alignment is shown in Figure 1 , which specifically includes the following steps:

[0019] S1. Pre-fine-tuning: First, initialize the target model using the pre-trained model parameters, and then perform N times of fine-tuning of the batch normalization layer parameters using the data of the target domain with L1 regularization, where N is usually set to 10. This step can effectively avoid pruning channels that are useful for the target domain data and affect subsequent fine-tuning.

[0020] S2. Pre-pruning: The absolute value of the batch normalization layer parameters of the trained target model represents the importance of the corresponding channel, and the last channel is pruned after global sorting, obtaining a pre-pruned model.

[0021] S3. Model pruning training based on weighted attention alignment: fine-tuning the pruned model, introducing the guidance of the pre-trained model through the weighted attention alignment mechanism, and obtaining a pruned model for the target domain dataset.

[0022] Since the batch normalization layer parameters are the basis for judging the importance of channels, they are widely used in model pruning tasks. Based on the weighted attention alignment mechanism, the parameters are extended to the pruning transfer task. On the one hand, the pruning and transfer are effectively combined, solving the problem that they are unrelated. On the other hand, after pruning the model channels, the number of channels is reduced, and the conventional feature alignment method in transfer learning cannot be applied. However, the attention mechanism can fuse the channel dimension, solving the feature alignment problem of inconsistent channel numbers.

[0023] In the weighted attention alignment mechanism, the calculation steps of the weighted attention map are as follows: obtaining the batch normalization layer parameters of the convolutional neural network where C represents the number of channels, and the intermediate tensor of the activation layer output after the batch normalization layer is obtained where H and W represent the height and width of the feature map, respectively;

[0024] The calculation of the weighted attention map is obtained using the function WAM(w, x), as shown in Figure 2 The specific steps include the following sub-steps:

[0025] S31. The batch normalization layer parameters w are normalized by softmax to obtain a weight vector

[0026] S32. The intermediate tensor is pooled where A represents the size of the pooled feature map;

[0027] S33. The weight vector is multiplied with the pooled intermediate tensor to obtain a weighted tensor

[0028] S34. The weighted tensor is summed by channel where represents the forward slice of the tensor ; then the weighted attention map is obtained by Frobenius normalization where F i,j represents the element in the i-th row and j-th column of the matrix.

[0029] In step S3, the pre-trained model is for a large dataset and cannot be changed, while the pruned model is for a small dataset and is generated by pruning the pre-trained model. A weighted attention alignment module, an output feature alignment module, and a target domain perception training module are set.

[0030] The weighted attention alignment module uses the weighted attention alignment mechanism to obtain the weighted attention maps of the pre-trained model and the pruned model, and aligns the weighted attention maps of the two at the corresponding positions, using the knowledge of the pre-trained model to guide the fine-tuning training of the pruned model to improve the pruning performance. The specific application location of the weighted attention alignment module is all the "convolutional layer-batch normalization layer-activation layer" network structure blocks in the network. For a network with a total of B such network structure blocks, its loss function is in denote the parameters of the batch normalization layer in the kth network structure block of the pre-trained model and the pruned model, respectively. denote the output feature tensors of the kth network structure block of the i-th image input pre-training model and pruned model, respectively. WAM(·,·) is the weighted attention map function. represents the square of the Frobenius norm.

[0031] Since the two models target different datasets, the output logical probability distribution sizes are different and cannot be aligned directly. Therefore, the output feature alignment module replaces the classification head of the pruned model with the classification head of the pre-trained model and then performs probability distribution alignment. Its loss function is in Denote the output probability distribution of the i-th picture after it is input into the feature extractor of the pre-trained model and the pruning model, and then input into the classification head of the pre-trained model, D KL (·) is the KL divergence function.

[0032] The target domain-aware training module fine-tunes the pruned model on a small dataset and uses data labels for supervised training. Its loss function is in represents the output probability distribution of the i-th image after it is input into the feature extractor of the pruning model and then input into the classification head of the pruning model, y i is the label of the i-th image, and CE(·) is the cross entropy function.

[0033] In summary, the loss function for a single image is:

[0034] Where α and β are balance coefficients. The pre-trained model is frozen, and the trainable model is the pruned model. Finally, the system obtains the pruned model for the small data set.

[0035] The model pruning system based on weighted attention alignment described in the present invention uses the model pruning method to prune the model to obtain a pruned model for a small data set.

[0036] The model pruning method and system based on weighted attention alignment described in the present invention have the following advantages:

[0037] (1) The weighted attention map calculation method combines pruning and transfer by leveraging the batch normalization layer parameters to quantify channel importance, thus increasing the interpretability of the pruning-transfer task. Furthermore, the channel fusion method improves the attention map, successfully solving the problem of feature alignment with inconsistent channel numbers and making it feasible to perform pruning and transfer tasks simultaneously.

[0038] (2) The weighted attention alignment training mechanism successfully realizes the joint training of pruning and transfer, solving the problem of high cost of step-by-step and iterative training in existing methods. It not only improves the performance of pruning by using pre-training knowledge, but also significantly reduces the time cost of the training process.

[0039] Those skilled in the art can make various other corresponding changes and deformations based on the technical solutions and concepts described above, and all of these changes and deformations should fall within the scope of protection of the claims of the present invention.

Claims

1. A model pruning method based on weighted attention alignment, characterized in that: The following steps are involved: S1. Pre-fine-tuning: Initialize the target model using the pre-trained model parameters, and then fine-tune the batch normalization layer parameters several times with L1 regularization using data from the target domain. S2. Pre-pruning: Globally sort the absolute values ​​of the batch normalization layer parameters of the trained target model and prune the last channels to obtain a pre-pruned model. S3. Model pruning training based on weighted attention alignment: Prune and fine-tune the pruned model. Using the weighted attention alignment mechanism, guidance from the pre-trained model is introduced to obtain a pruned model for the target domain dataset. In the weighted attention alignment mechanism, the calculation steps of the weighted attention map are as follows: Get the batch normalization layer parameters of the convolutional neural network Where C represents the number of channels, and the intermediate tensor output by the activation layer after the batch normalization layer is obtained Where H and W represent the height and width of the feature map respectively; The weighted attention map is calculated using the function Obtaining, including the following sub-steps: S31. Perform softmax normalization on the batch normalization layer parameter w to obtain the weight vector S32. Pooling the intermediate tensor Where A represents the size of the feature map after pooling; S33. The weight vector and the intermediate tensor after pooling Multiply by channel to get weighted tensor S34. Weighted tensor channel summation in Represents the forward slice of the tensor; then the weighted attention map is obtained after Frobenius normalization Among them F i,j Represents the element in the i-th row and j-th column of the matrix; In step S3, a weighted attention alignment module, an output feature alignment module, and a target domain perception training module are set; The weighted attention alignment module uses the weighted attention alignment mechanism to obtain the weighted attention maps of the pre-trained model and the pruned model, and aligns the weighted attention maps of the two at the corresponding positions, and uses the knowledge of the pre-trained model to guide the fine-tuning training of the pruned model to improve the pruning performance; for a network with a total of B network structure blocks, the loss function is in denote the parameters of the batch normalization layer in the kth network structure block of the pre-trained model and the pruned model, respectively. denote the output feature tensors of the kth network structure block of the i-th image input pre-training model and pruned model, respectively. WAM(·,·) is the weighted attention map function. represents the square of the Frobenius norm; The output feature alignment module replaces the classification head of the pruned model with the classification head of the pre-trained model, and then performs probability distribution alignment; the loss function is in Denote the output probability distribution of the i-th picture after it is input into the feature extractor of the pre-trained model and the pruning model, and then input into the classification head of the pre-trained model, D KL (·) is the KL divergence function; The target domain-aware training module fine-tunes the pruned model on a small dataset and uses data labels for supervised training. The loss function is in represents the output probability distribution of the i-th image after it is input into the feature extractor of the pruning model and then input into the classification head of the pruning model, y i is the label of the i-th image, CE(·) is the cross entropy function; The loss function for a single image is: Among them, α and β are balance coefficients; Finally, we get a pruned model for small datasets.

2. The model pruning system based on weighted attention alignment according to claim 1, characterized in that The model is pruned using the model pruning method as described in claim 1 to obtain a pruned model for a small data set.

Images