Image classification system and method based on lightweight LA Transformer network
By introducing a local attention mechanism and a convolutional neural network into the visual Transformer architecture, a lightweight LA Transformer architecture is constructed, which solves the problem of high complexity in the visual Transformer architecture and achieves more efficient image classification performance.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- TIANJIN UNIV
- Filing Date
- 2023-01-09
- Publication Date
- 2026-05-29
AI Technical Summary
Existing visual Transformer architectures suffer from problems such as high training difficulty, high model complexity, and high computational cost in image classification tasks. Furthermore, existing lightweight solutions have failed to effectively address the excessive time and space complexity of self-attention modules and feedforward neural networks.
We adopt a lightweight LA Transformer architecture based on local attention mechanism. Local feature extraction is performed through a convolutional neural subnetwork with residual connections. It combines a local self-attention subnetwork and an attention feedforward neural subnetwork to replace the traditional self-attention module and feedforward neural network, reducing the number of parameters and computational cost while retaining the feature extraction capability.
While reducing the number of parameters and computational cost, it improves image classification accuracy and achieves more efficient feature extraction and classification performance, especially performing well on the ImageNet-1k dataset.
Smart Images
Figure CN115984620B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the fields of image recognition technology and machine learning technology, and specifically to an image classification method based on a lightweight visual Transformer architecture designed with a channel attention mechanism. Background Technology
[0002] Image classification is a crucial problem in modern computer vision. The performance and efficiency of deep neural network models in image classification often directly reflect their feature extraction capabilities. In recent years, the excellent performance of Transformers in natural language processing has attracted significant attention from computer vision researchers. Compared to traditional convolutional neural networks, the visual Transformer architecture has a larger receptive field, capturing global information from the image data at the network's input. Therefore, the visual Transformer architecture exhibits stronger feature extraction capabilities and demonstrates higher performance in tasks such as image classification and object detection. However, mainstream visual Transformer architectures still suffer from drawbacks such as high training difficulty and higher model complexity. In the current field of computer vision, lightweight Transformer architectures are a highly focused research area. Researchers hope to improve the architecture of visual Transformers to reduce model time and space complexity without compromising feature extraction capabilities, achieving a better balance between efficiency and performance.
[0003] The mainstream visual Transformer architecture typically consists of multiple isomorphic Transformer blocks stacked together. Each Transformer block includes a self-attention module for modeling spatial correlations of feature maps and a feedforward neural network for modeling inter-channel correlations. Both the self-attention module and the feedforward neural network utilize numerous fully connected layers to process features, introducing a large number of learnable parameters and significantly increasing the model's space complexity. Furthermore, the time complexity of the self-attention module during feature extraction from image data is proportional to the square of the image size, making it difficult to control the computational load of traditional visual Transformer models when processing high-resolution images. Therefore, while the visual Transformer architecture possesses strong feature extraction capabilities, it also introduces significant time and space complexity.
[0004] To achieve lightweighting of the visual Transformer architecture, there are currently two main technical approaches:
[0005] The first approach replaces the token embedding part in the visual Transformer architecture with a small convolutional neural network. Based on the inherent local correlation of image data, this solution addresses the visual Transformer architecture's inability to effectively capture local spatial information. Lightweight Transformer architectures built using this design often achieve similar feature extraction capabilities to the original architecture with fewer stacked Transformer blocks, thus achieving a lighter model. However, this approach does not fundamentally solve the problem of excessive time and space complexity in the self-attention module and feedforward neural network of the traditional Transformer architecture.
[0006] The second approach involves dynamic pruning of feature maps in the Transformer architecture. This approach primarily addresses the issue of excessive time complexity in modeling the spatial correlation of features by the self-attention module. Several works within this approach have proposed their own dynamic patching and pruning strategies, which can fuse feature tokens with similar semantic information to reduce the time complexity in computing self-attention. However, current works still have shortcomings. For example, they cannot effectively address the large number of parameters in the visual Transformer architecture, and in practical applications, this approach still incurs a significant storage burden on the model. Furthermore, the proposed token fusion schemes are not efficient enough, failing to guarantee that the model's feature extraction capability will not be significantly affected when fusing a certain proportion of tokens.
[0007] Therefore, improving the model architecture to achieve more efficient modeling of image features remains a significant challenge. Summary of the Invention
[0008] To overcome the shortcomings of existing methods and technologies, this invention aims to propose an image classification method based on a lightweight LA Transformer architecture. A novel lightweight visual Transformer architecture, LATransformer (Local Attention Transformer), based on a local attention mechanism, is constructed for the recognition of different features in image information, thereby achieving image classification.
[0009] This invention is achieved using the following technical solution:
[0010] An image classification system based on a lightweight LA Transformer network, comprising a feature extraction module, an LA Transformer main network, and an image classification module, wherein:
[0011] The feature extraction module is responsible for extracting local features from the input RGB image using a convolutional neural network with residual connections, thereby achieving downsampling of the input image and outputting... As input to the main LA Transformer network; the convolutional neural subnetwork with residual connections further includes two two-dimensional convolutional layers, a global max pooling layer, and a batch normalization layer structure, wherein the batch normalization layer structure includes three cascaded batch normalization layer substructures;
[0012] The LA Transformer main network is responsible for further feature extraction of image features using multiple cascaded LA Transformer blocks. The LA Transformer main network includes multiple cascaded LA Transformer blocks. Each LA Transformer block consists of a local self-attention subnetwork and a normalization operation layer, and a normalization operation layer and an attention feedforward neural subnetwork. The local self-attention subnetwork uses the local self-attention mechanism to model the spatial correlation information between patches, and the attention feedforward neural subnetwork models the correlation information between channels.
[0013] The image classification processing module is responsible for generating the probability for each image classification and completing the classification of the input RGB image.
[0014] An image classification method based on a lightweight LA Transformer network, comprising the following steps:
[0015] Step S1: A convolutional neural subnetwork with residual connections is used to extract local features from the input RGB image, achieving downsampling of the input image and outputting... As input to the main network of the LA Transformer;
[0016] Step S2: Further feature extraction of the image features is performed using multiple cascaded LA Transformer blocks. The main LA Transformer network includes multiple cascaded LA Transformer blocks. Each LA Transformer block consists of a local self-attention subnetwork and a normalization operation layer, and a normalization operation layer and an attention feedforward neural subnetwork. The local self-attention subnetwork uses the local self-attention mechanism to model the spatial correlation information between patches, and the attention feedforward neural subnetwork models the correlation information between channels.
[0017] Step S3: Generate the probability for each image classification to complete the classification of the input RGB image.
[0018] Compared to the classic ViT model, the image classification system and method based on the lightweight LA Transformer network of this invention can achieve higher image classification accuracy while having lower parameter and computational costs. The results can be verified on public datasets such as ImageNet-1k. Attached Figure Description
[0019] Figure 1 This is a structural diagram of the image classification system based on the lightweight LA Transformer network of the present invention;
[0020] Figure 2 Here is a schematic diagram of the architecture of a convolutional neural network (RNB) with residual connections; (2a) shows the overall structure of the RNB, and (2b) shows the Bottleneck substructure of the RNB.
[0021] Figure 3 A diagram of the main network structure for each LA Transformer;
[0022] Figure 4 This is a structural diagram of a Local Self-Attention Module (LSA).
[0023] Figure 5 This is a diagram of the Attention Feedforward Neural Network (AFFN) structure.
[0024] Figure 6 This is a flowchart of the image classification method based on the lightweight LA Transformer network of the present invention. Detailed Implementation
[0025] The technical solution will now be described in detail with reference to the accompanying drawings and embodiments. All other embodiments and technical substitutions of the embodiments obtained by those skilled in the art based on the embodiments of the present invention without departing from the spirit of the invention and without creative effort will fall within the scope of protection of the present invention.
[0026] like Figure 1 The diagram shows the structure of the image classification system based on the lightweight LA Transformer network of the present invention. The system includes a feature extraction module 100, an LA Transformer main network 200, and an image classification processing module 300.
[0027] The feature extraction module 100 uses a convolutional neural network (ResNet block, RNB) with residual connections to extract the RGB images of the input image data. Preliminary feature extraction is performed, resulting in the original feature map. As input; where N and C are a set of hyperparameters used in the visual Transformer architecture design, representing the number of patch feature maps in the Transformer block and the number of channels in each patch feature map, respectively, and W in and H in These represent the width and height of the input image, respectively.
[0028] like Figure 2 The diagram shows a schematic of a convolutional neural network (RNB) with residual connections. (2a) shows the overall structure of the RNB, which includes a two-dimensional convolutional layer Conv2d at the input of the subnetwork, a global max pooling layer, a Bottleneck structure, and a two-dimensional convolutional layer Conv2d at the input of the subnetwork. The Bottleneck structure includes three cascaded Bottleneck substructures. Each Bottleneck substructure includes two heterogeneous branch structures. The left branch structure consists of three two-dimensional convolutional layers and a batch normalization layer stacked alternately. The kernel size of the first and third two-dimensional convolutional layers is 1×1, while the kernel size of the second two-dimensional convolutional layer is 3×3, and the stride of the convolution is 2×2, which downsamples the input features. For input feature maps of any size, the number of channels of the output features of the first and second two-dimensional convolutional layers is 64, while the last two-dimensional convolutional layer performs a dimensionality increase operation to increase the number of channels to 256. Finally, the modified linear unit activation function (ReLU) is used to perform non-linear activation on the features, which is then used as the output of the left branch structure. The right branch structure consists of only one 2D convolutional layer with a kernel size of 1×1 and a stride of 2×2, and a batch normalization layer, with 256 output feature channels, the same as the output feature channels of the left branch. Both branches undergo a downsampling operation. The image feature size is processed to (64, 112, 112) pixels. The output features of the two branches are superimposed as the output features of the Bottleneck substructure. After processing through three cascaded Bottleneck substructures, a 2D convolutional layer (ConV2d) with a kernel size of 1×1 is used to reduce the number of feature channels to 192, and the feature size is downsampled to 14×14, which is then used as the output feature map of the RNB. The output feature map is stretched in spatial dimension to obtain the final output features of the transformed RNB. N=196, C=192, which serve as the inputs to the LATransformer main network 200.
[0029] The LA Transformer main network 200 consists of multiple stacked LA Transformer blocks, each of which further includes a Local Self-Attention (LSA) subnetwork and an Attention Feed-Forward (AFFN) subnetwork. The LSA subnetwork models the spatial correlation information between image feature patches using a local self-attention mechanism, while the AFFN models the correlation information between image feature channels. Specifically, based on image features, convolutional, fully connected, or self-attention operations are used in the network, employing learnable parameters to capture correlation information between channels, between pixels in space, or between patches in space. Therefore, the forward propagation process of image features here is the process of modeling correlation information.
[0030] like Figure 3 The diagram shows the structure of each LA Transformer block. The overall architecture of an LA Transformer block consists of two parts: a Local Self-Attention Subnetwork (LSA) and a Normalization Operation Layer (LA Transformer), followed by an Attention Feedforward Neural Network (AFFN) within the LA Transformer block. Compared to the classic Transformer block architecture, the innovation lies in replacing the global self-attention network in the traditional Transformer block with the Local Self-Attention Subnetwork (LSA), and replacing the feedforward neural network in the traditional Transformer block with the Attention Feedforward Neural Network.
[0031] like Figure 4 The diagram shown is a structure diagram of a Local Self-Attention Subnetwork (LSA).
[0032] In the local self-attention subnetwork, the feature map is first... The token dimension N×C of each feature is restored to C×H×W, which is represented as a feature map. Feature maps were analyzed along the C, H, and W dimensions respectively. Pooling operations are performed on each feature to obtain X0_c is the pooled feature map in dimension C, X0_h is the pooled feature map in dimension H, and X0_w is the pooled feature map in dimension W. For the three vectors X0_c, X0_w, and X0_w, convolution operations are performed using three one-dimensional convolutional layers with a kernel size of 3. During the process of the one-dimensional convolution operator traversing the image features in the three dimensions H, W, and C, the learnable parameters in the three Conv1d convolutional kernels will spontaneously model the correlation between adjacent channels or adjacent patches.
[0033] Using the Sigmoid activation function for non-linear activation processing, three feature maps Y_c, Y_h, and Y_w are obtained, as follows:
[0034] Y0_c=X0_c*sigmoid(Conv1d(avgpool c (X0)))
[0035] Y0_h=X0_h*sigmoid(Conv1d(avgpool H (X0)))
[0036] Y0_w=X0_w*sigmoid(Conv1d(avgpool W (X0)))
[0037] The feature maps Y0_c, Y0_h, Y0_w ∈ R are obtained by performing a dot product of the three attention weight vectors and the feature maps Y0_c, Y0_h, Y0_w in their respective dimensions. C×H×W These three features are concatenated along the channel dimension, and the concatenated feature map is reduced in dimensionality using a 1×1 convolutional layer to obtain the final output feature map. And in subsequent operations, the size was changed to
[0038] like Figure 5 The diagram shown is a structural diagram of the Attention Feedforward Neural Network (AFFN). For the feature map input to the AFFN... It is divided into two equal parts along the channel dimension, as shown below:
[0039]
[0040]
[0041]
[0042] Enter the two branch structures respectively:
[0043] The attention (att) branch on the left consists of an average pooling layer (Avgpool), two fully connected layers (FClayer), and a sigmoid activation function. The input feature map for this branch... To compress information in the spatial dimension using GAP (Global Average Pooling), the expression is as follows:
[0044]
[0045] Then, by using two fully connected layers and the Sigmoid activation function, channel attention weights W2 and W1 are obtained. These channel attention weights are then used to reweight the image features output from the other branch along the channel dimension.
[0046] After that, regarding y att The output feature map Y formed by the set of att ∈R 1×(C / 2) It consists of two fully connected (FC) layers and a sigmoid activation function, and the calculation formula is as follows:
[0047] z = W2ReLU(W 1y ),z∈R 1×C / 2
[0048] Here, two fully connected (FC) layers are used to model the correlation between channels based on the global spatial information of each channel, and to generate channel attention weights z.
[0049] The feedforward neural network branch structure (FFN) on the right includes two fully connected layers (FC layers), and the input feature map for this branch... The process of first increasing dimensionality and then reducing dimensionality is expressed as follows:
[0050]
[0051] The output feature map expression is as follows:
[0052]
[0053] This enables the modeling of feature maps along the channel dimension;
[0054] Finally, the channel attention weight z from the output of the left attention (att) branch is used to evaluate the feature map output from the right branch structure. Perform channel-wise dot product and reweight along the channel dimension to output the final output feature map Y of AFFN, as shown in the following expression:
[0055]
[0056] The mathematical expressions used in the above process can be summarized as follows:
[0057] The feature maps initially extracted by RNB are represented as follows:
[0058] X0 = f RNB (X in (1)
[0059] The feature maps corresponding to each LA Transformer block are represented as follows:
[0060]
[0061] Where t is the number of the LA Transformer block.
[0062] The LA Transformer blocks from the first to the t-th have the same structure, but as the network extracts features layer by layer, the LA Transformer blocks closer to the output spontaneously learn higher-level semantic features. The output of the t-th LA Transformer block is then used...
[0063] The image classification module 300 uses the output feature map of the attention feedforward neural subnetwork. As input, global average pooling is performed, and then a fully connected layer (FC layer) and a softmax function are used to generate the probability for each image classification, thus completing the classification of the input image data. This module follows the classic design of the visual Transformer architecture.
[0064] The LA Transformer architecture proposed in this invention innovates from three main perspectives. Here, we first provide a qualitative analysis of these three main innovations:
[0065] First, a convolutional neural network (RNB) with residual connections is used to extract local features from image data, effectively improving the model's performance on image classification tasks while introducing fewer additional parameters. The design of this part of the structure mainly references the classic convolutional neural network architecture ResNet, which is more conducive to model training. This addresses the limitation of the visual Transformer architecture in effectively capturing local information in space. The lightweight Transformer architecture built using this design approach can often achieve similar feature extraction capabilities to the original architecture with fewer stacked Transformer blocks, thus achieving a lightweight effect on the model. Furthermore, it fundamentally solves the problem of excessive time and space complexity of the self-attention module and feedforward neural network in the traditional Transformer architecture.
[0066] Secondly, replacing the self-attention module in the traditional visual Transformer architecture with a local self-attention submodule introduces negligible additional parameters while retaining basic local spatial feature modeling capabilities. Furthermore, this local attention module offers a significant advantage in computational cost when dealing with large-sized input image features. It effectively models the spatial information of feature maps while substantially reducing the number of parameters and computational cost. The attention-based feedforward neural network introduces a channel attention mechanism into the traditional feedforward neural network, further improving the network's ability to model the correlations between feature channels.
[0067] Finally, the third innovative attention-forward neural network utilizes the important role of the channel attention module in fusing information between channels. While reducing the number of channels that the feedforward neural network processes features, it can maintain the ability to fuse information between channels, effectively improving the model's feature extraction capability.
[0068] Specific implementation: The ImageNet dataset is used for image classification tasks. When training the network model on the ImageNet dataset, this invention adopts the same data augmentation methods as ViT, including size, color, and flip enhancements, as well as mixup, randAugment, cutmix, and label smoothing schemes.
[0069] In the training phase of the image classification task, this example uses the training set from the ImageNet-1K dataset to train the LATransformer architecture. The AdamW optimization algorithm and cosine annealing training strategy are employed during training. The weight decay parameter is set to 1e-4, the momentum parameter to 0.9, the batch size of the training data to 256, the initial learning rate to 0.1, and the model training iterations to 300 epochs.
Claims
1. An image classification system based on a lightweight LA Transformer network, characterized in that, The system includes a feature extraction module, an LA Transformer main network, and an image classification module, wherein: The feature extraction module is responsible for extracting local features from the input RGB image using a convolutional neural network with residual connections, thereby achieving downsampling of the input image and outputting... As the input to the main network of the LA Transformer, N=196, representing the number of patch feature maps; C=192, representing the number of channels in each patch feature map; the convolutional neural subnetwork with residual connections further includes a two-dimensional convolutional layer, a global max pooling layer, a Bottleneck structure at the input end of the subnetwork, and a two-dimensional convolutional layer at the output end of the subnetwork, wherein the Bottleneck structure includes three cascaded Bottleneck substructures; The LA Transformer main network is responsible for further feature extraction of image features using multiple cascaded LA Transformer blocks. The LA Transformer main network includes multiple cascaded LA Transformer blocks. Each LA Transformer block consists of two parts: a Local Self-Attention Subnetwork (LSA) and a normalization operation layer, and a normalization operation layer and an Attention Feedforward Neural Network (AFFN). The LSA uses the local self-attention mechanism to model the spatial correlation information between image feature patches and patches, and the AFFN uses the attention feedforward neural network to model the correlation information between image feature channels. The image classification module is responsible for generating the probability for each image category, thus completing the classification of the input RGB image.
2. The image classification system based on a lightweight LA Transformer network as described in claim 1, characterized in that, Each Bottleneck substructure further includes two branch structures. One branch structure consists of three alternating stacked 2D convolutional layers and a batch normalization layer. The first and third 2D convolutional layers have a kernel size of 1×1, the second 2D convolutional layer has a kernel size of 3×3, and the convolution stride is 2×2. Finally, the features are non-linearly activated using the modified linear unit activation function ReLU. The other branch structure includes a 2D convolutional layer with a kernel size of 1×1 and a convolution stride of 2×2, and a batch normalization layer, which performs a downsampling operation. The output features of the two branch structures are superimposed as the output features of each Bottleneck substructure.
3. The image classification system based on a lightweight LA Transformer network as described in claim 1, characterized in that, The local self-attention subnetwork includes an average pooling layer, a one-dimensional convolutional layer, and a sigmoid function set in each channel dimension C, H, W.
4. The image classification system based on a lightweight LA Transformer network as described in claim 1, characterized in that, The attention feedforward neural subnetwork is divided into two parts in the channel dimension. The left branch structure includes an average pooling layer, two fully connected layers and a sigmoid activation function; the right branch structure includes two fully connected layers.
5. The image classification system based on a lightweight LA Transformer network as described in claim 1, characterized in that, In the local self-attention subnetwork, the feature map is first... The token dimension N×C of each feature is restored to C×H×W, which is represented as a feature map. The feature maps were analyzed along the C, H, and W dimensions, respectively. Pooling operations are performed on each feature to obtain , The pooled feature map on dimension C, Pooled features on dimension H For the pooled feature map on dimension W, given three vectors... , , Convolution operations are performed using three one-dimensional convolutional layers with a kernel size of 3. During the process of the one-dimensional convolution operator traversing the image features in the three dimensions of H, W, and C, the learnable parameters in the three Conv1d convolutional kernels will spontaneously model the correlation between adjacent channels or adjacent patches. Three feature maps were obtained by performing non-linear activation using the Sigmoid activation function. , , , means as follows: Using three attention weight vectors and feature maps , , The feature map is obtained by performing a dot product in their respective dimensions. These three features are concatenated along the channel dimension, and the concatenated feature map is reduced in dimensionality using a 1×1 convolutional layer. After resizing, the final output feature map is obtained. .
6. The image classification system based on a lightweight LA Transformer network as described in claim 1, characterized in that, In the attention-feedforward neural subnetwork, for the feature map of the input AFFN It is divided into two parts on an average of two dimensions along the channel dimension, as shown below: Enter the two branch structures respectively: The attention branch on the left consists of an average pooling layer (Avgpool), two fully connected layers (FC), and a sigmoid activation function. The input feature map for this branch... The information in the spatial dimension is compressed using GAP global average pooling, as shown in the following expression: Then, by using two fully connected layers and a sigmoid activation function, the channel attention weights are obtained. , Channel attention weights are used to reweight the image features output from the other branch along the channel dimension. H , W These are the dimensions of each channel; for The output feature map formed by the set of Use two FC layers and Sigmoid activation. The function is composed of the following formulas: Here, two fully connected (FC) layers are used to model the correlation between channels based on the global spatial information of each channel, and to generate channel attention weights. z ; The feedforward neural network branch structure FFN on the right includes two fully connected layers (FC), and the input feature map for this branch... The process involves first increasing the dimensionality and then decreasing it. The expression for the processing flow is as follows: The output feature map expression is as follows: This enables the modeling of feature maps along the channel dimension; Finally, the channel attention weight z from the output of the left attention branch (att) is used to evaluate the feature map output from the right branch structure. Perform channel-wise dot product and reweight along the channel dimension to output the final feature map of the Attention Feedforward Neural Network (AFFN). Y The expression is as follows: 。 7. An image classification method based on a lightweight LA Transformer network, characterized in that, The method includes the following steps: Step S1: A convolutional neural subnetwork with residual connections is used to extract local features from the input RGB image, achieving downsampling of the input image and outputting... As the input to the main network of the LA Transformer, N=196, representing the number of patch feature maps; C=192, representing the number of channels in each patch feature map; the convolutional neural subnetwork with residual connections further includes a two-dimensional convolutional layer, a global max pooling layer, a Bottleneck structure at the input end of the subnetwork, and a two-dimensional convolutional layer at the output end of the subnetwork, wherein the Bottleneck structure includes three cascaded Bottleneck substructures; Step S2: Further feature extraction of image features is performed using multiple cascaded LA Transformer blocks. The main LA Transformer network includes multiple cascaded LA Transformer blocks. Each LA Transformer block consists of two parts: a Local Self-Attention Subnetwork (LSA) and a normalization operation layer, and a normalization operation layer and an Attention Feedforward Neural Network (AFFN). The LSA is used to model the spatial correlation information between image feature patches and patches using the local self-attention mechanism, and the AFFN is used to model the correlation information between image feature channels. Step S3: Using the feature maps extracted in step S2, generate the probability for each image classification to complete the classification of the input RGB image.