Image classification algorithm based on PSA-Swin Transformer
By introducing PSA-Swin Transformer, EPE module and PSA module into the image classification algorithm, the problem of poor performance in the existing technology on small-scale data sets is solved, and the effect of efficient image classification under small data volume is achieved.
Patent Information
- Application Number
- CN202410793055.6
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-06-19
- Publication Date
- 2025-02-18
AI Technical Summary
Existing deep learning-based image classification methods do not perform well on small-scale data sets, and rely on large-scale data sets and pre-training weights, making it difficult to effectively process smaller data volumes.
The image classification algorithm based on PSA-Swin Transformer is adopted. The algorithm blocks and flattenss the pictures through the Patch Partition module, combines the EPE module and the PSA module to build feature maps of different sizes, and finally through Layer Norm, global pooling layer and full connection The layer gets the final output.
This algorithm can effectively process small data sets, improve performance on small-scale data sets, does not rely on pre-training weights, and enhances the model's adaptability to small data sets.
Smart Images
Figure CN119478476A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the field of deep learning image classification, and specifically to an image classification algorithm based on PSA-Swin Transformer. Background Art
[0002] Image classification refers to the process of identifying the category of a given image. It is a basic task in the field of computer vision. Its goal is to assign the input image to a predefined category. This task involves extracting features from the image and then mapping these features to predefined categories using a classification algorithm. Image classification has a wide range of uses in many practical applications, such as face recognition, handwritten digit recognition, vehicle recognition, etc.
[0003] At present, most of the existing deep learning-based image classification methods have the limitation of being heavily dependent on large-scale datasets and extensive pre-training. The existing deep learning-based image classification methods can be roughly divided into CNN-based methods and Transformer-based methods. Among them, CNN uses spatial inductive bias to effectively learn image representations, has fewer parameters and can efficiently learn sample features, but its performance may be limited; Transformer can learn global representations and performs well in tasks such as image classification, but usually requires large-scale datasets and pre-training, and its generalization ability may not be as good as CNN. For example, visual transformers (ViTs) will show good performance after being pre-trained on large datasets such as ImageNet-1k / 22k or JFT-300M datasets, but the effect may not be good for small-scale datasets.
[0004] In practice, collecting high-quality labeled data or human annotators is very expensive. Synthetic data, while promising, is often a distorted version of real data, and any modeling or reasoning done on synthetic data carries additional risks. Models trained on synthetic data often need to be fine-tuned with real data before deployment. In addition, methods such as transfer learning often fail to fully solve the problem because the pre-training datasets are biased and do not reflect the environment. Summary of the invention
[0005] The purpose of the present invention is to solve the problems existing in the prior art and provide an image classification algorithm based on PSA-SwinTransformer, which is a method suitable for processing small data sets and can effectively process smaller data volumes without the need for pre-training weights.
[0006] In order to achieve the above object, the present invention is implemented through the following technical solutions:
[0007] An image classification algorithm based on PSA-Swin Transformer, comprising the steps of:
[0008] S1. First, the image is input into the Patch Partition module for block division, and then flattened in the channel direction;
[0009] S2, construct feature maps of different sizes through four stages in sequence;
[0010] In the first stage, the data first passes through a Linear Embedding layer and then passes through a PSA-SwinTransformer Block layer.
[0011] In the second, third, and fourth stages, the image is first downsampled through a Patch Merging layer and then through a PSA-Swin Transformer Block layer;
[0012] The PSA-Swin Transformer Block layer includes the EPE module and the PSA module;
[0013] S3, then connect a Layer Norm layer, a global pooling layer and a fully connected layer to get the final output.
[0014] Preferably, in step S1, the image is an RGB three-channel image, each patch has 4×4=16 pixels, each pixel has three values of R, G, and B, and after flattening, it is 16×3=48. After passing through the Patch Partition layer, the shape of the image changes from [H, W, 3] to [H / 4, W / 4, 48]; then, a linear transformation is performed on the channel data of each pixel through the Linear Embeding layer, that is, the original number of channels is changed from 48 to C, and the shape of the image is changed from [H / 4, W / 4, 48] to [H / 4, W / 4, C], where H represents the height of the image, W represents the width, and C represents the number of channels.
[0015] Preferably, in step S2, when the feature map passes through the first Patch Merging layer, the Patch Merging layer divides each 2x2 adjacent pixels into a patch, and then splices the pixels at the same position in each patch together to obtain 4 times the number of original feature maps; then the feature map is spliced in the channel direction, and then passes through a layer of normalization layer, and finally passes through a fully connected layer to make a linear change in the depth direction of the feature map, so that the depth of the feature map is changed from C at the input to 2C, and the resolution is halved; after the second stage, the output resolution is H / 8×W / 8, and the number of channels becomes 2C; after the third stage, the output resolution is H / 16×W / 16, and the number of channels becomes 4C; after the fourth stage, the output resolution is H / 32×W / 32, and the number of channels becomes 8C.
[0016] Preferably, in step S2, EPE(X)=GSConv(X), wherein GSConv uses a shuffle operation to allow information exchange between groups; GSConv first distributes input features to different groups through a grouped convolution architecture, and independently performs convolution operations within each group.
[0017] Preferably, in step S2, the PSA module accurately implements pixel-by-pixel regression by combining two key strategies:
[0018] First, the polarization filtering strategy completely folds the input tensor along the corresponding dimension while maintaining the internal resolution when performing channel and spatial attention calculations;
[0019] Second, the enhancement strategy is achieved by introducing nonlinearity in the attention mechanism.
[0020] Preferably, in the PSA module, for the input feature X∈R C×H×W ,
[0021] When only channel attention is used:
[0022]
[0023] Among them, Wq, Wk, and Wv are 1×1 convolutional layers;
[0024] σ1, σ2 are reshape operators, F SM is the softmax operator, and × is the matrix dot product operation
[0025] The number of internal channels between Wv|Wq and Wz is C / 2;
[0026] The output of the channel branch only is: Among them, ch is the channel multiplication operator;
[0027] When only spatial attention is used:
[0028] A sp (X) = F SG [σ3(F SM (σ1(F GP (W q (X))))×σ2(W v (X)))]
[0029] Wq, Wv are standard 1×1 convolutions, and θ2 is the intermediate parameter of these channel convolutions;
[0030] σ1σ2σ3 are three reshape operators, F SM is the softmax operator, F GP (·) is the global pooling operator;
[0031] The output of the space-only branch is: Among them, sp is the spatial multiplication operator;
[0032] The outputs of the above two branches are composed in a parallel layout:
[0033] PSA p (X) = Z ch +Z sp =A ch (X)⊙ ch X+A sp (X)⊙ sp X.
[0034] Compared with the prior art, the beneficial effects of the present invention are:
[0035] 1. This paper adopts an efficient position embedding module (EPE), which aims to greatly improve the parameter utilization efficiency and network expressiveness through a grouped convolutional architecture, that is, first assigning input features to different groups, and then performing convolution operations independently in each group. The special feature of EPE is its shuffle operation, which allows information to be exchanged between groups, thereby enhancing the network's ability to integrate different features and more effectively capture global and local feature information. Information exchange enhances the network's ability to integrate different features, thereby more effectively capturing global and local feature information.
[0036] 2. The attention mechanism in deep convolutional neural networks can improve the long-range dependency of the model, but the learning of specific element attention (such as non-local blocks) is complex and sensitive to noise. Most simplified attention mixtures attempt to achieve the best compromise in multiple types of tasks, so they cannot fully exhaust their representation capabilities in pure channel and pure spatial branches. To this end, the present invention introduces a polarized self-attention module, which can effectively solve the above problems by combining the two key designs of polarized filtering and augmentation. BRIEF DESCRIPTION OF THE DRAWINGS
[0037] Figure 1 It is the overall flow chart of the present invention;
[0038] Figure 2 is a schematic diagram of the EPE module;
[0039] Figure 3 is a schematic diagram of a PSA module. DETAILED DESCRIPTION
[0040] The present invention will be further described below in conjunction with specific embodiments. It should be understood that these embodiments are only used to illustrate the present invention and are not intended to limit the scope of the present invention. In addition, it should be understood that after reading the content taught by the present invention, those skilled in the art can make various changes or modifications to the present invention, and these equivalent forms fall within the scope limited by the application equally.
[0041] Embodiment 1: As attached Figure 1-3 As shown, the present invention is an image classification algorithm based on PSA-Swin Transformer, comprising the steps of:
[0042] S1. First, input the image into the Patch Partition module for block division, that is, each 4×4 adjacent pixels is a patch, and then flatten it in the channel direction. Assuming that the input is an RGB three-channel image, then each patch has 4×4=16 pixels, and each pixel has three values of R, G, and B, so after flattening it is 16×3=48. Therefore, after passing through the PatchPartition layer, the shape of the image changes from [H, W, 3] to [H / 4, W / 4, 48]. Then, through the Linear Embedding layer, a linear transformation is performed on the channel data of each pixel, that is, the original number of channels changes from 48 to C, and the shape of the image changes from [H / 4, W / 4, 48] to [H / 4, W / 4, C]. Among them, H represents the height of the image, W represents the width, and C represents the number of channels.
[0043] S2, construct feature maps of different sizes through four stages in sequence;
[0044] In the first stage, the data first passes through a Linear Embedding layer and then passes through a PSA-SwinTransformer Block layer.
[0045] In the second, third, and fourth stages, the image is first downsampled through a Patch Merging layer and then through a PSA-Swin Transformer Block layer;
[0046] The PSA-Swin Transformer Block layer includes the EPE module and the PSA module;
[0047] Except for the first stage where a Linear Embedding layer is passed, the remaining three stages all first pass through a Patch Merging layer for downsampling to reduce the number of tags. When the feature map passes through the first PatchMerging layer, the Patch Merging layer divides each 2x2 adjacent pixel into a patch, and then splices the pixels at the same position in each patch together to obtain 4 times the number of original feature maps. The feature map is then spliced in the channel direction, passed through a normalization layer, and finally passed through a fully connected layer to make a linear change in the depth direction of the feature map, changing the depth of the feature map from C at the input to 2C, while the resolution is halved. That is, after passing through the PatchMerging layer, the height and width of the feature map will be halved, and the depth will be doubled. After the second, third, and fourth stages, the output resolution is H / 8×W / 8, H / 16×W / 16, H / 32×W / 32, and the number of channels will become 2C, 4C, and 8C.
[0048] S3, then connect a Layer Norm layer, a global pooling layer and a fully connected layer to get the final output.
[0049] Among them: Efficient Position Embedding Module (EPE)
[0050] Previous methods mainly use absolute or relative position embedding of image tasks to solve the problem of position encoding. In addition, due to the lack of absolute position information, the relative version will change the self-attention and perform poorly. In order to overcome the above problems, the present invention extends the conditional position embedding and designs an Efficient position embedding (EPE) module. The EPE module is as follows Figure 2 shown.
[0051] EPE(X) = GSConv(X), where GSConv uses a grouped convolution architecture to first distribute input features to different groups, and then independently performs convolution operations within each group. This approach greatly improves parameter utilization efficiency and network expressiveness. The special feature of GSConv is its shuffle operation, which allows information exchange between groups and enhances the network's ability to integrate different features, thereby more effectively capturing global and local feature information.
[0052] In image classification tasks, the spatial relationship of images is particularly important. Through its grouping and shuffling strategies, GSConv can more accurately capture the spatial patterns and contextual relationships in images, which is crucial for recognizing and classifying complex patterns in images.
[0053] GSConv not only improves computational efficiency, but its network enhanced by feature mixing and information exchange can provide richer and more detailed feature representation, which is extremely important in image classification tasks and helps improve the model's understanding of image content and classification accuracy. Therefore, in scenarios that require accurate processing and recognition of complex spatial relationships in images, EPE shows better performance than CPE and DPE.
[0054] Polarized Self-Attention Module (PSA)
[0055] The polarized self-attention module achieves accurate pixel-wise regression by combining two key strategies:
[0056] First, the polarization filtering strategy completely folds the input tensor along the corresponding dimension while maintaining the internal resolution when performing channel and spatial attention calculations, ensuring a small number of parameters and high resolution.
[0057] Secondly, the enhancement strategy can make the model output more refined and closer to the real situation by introducing nonlinearity in the attention mechanism (such as using 2D Gaussian distribution for keypoint heatmap and 2D Binomial distribution for binary segmentation mask), thus making the model output more refined and realistic.
[0058] Figure 3 The PSA module is described, which preserves channel information as much as possible in channel attention and spatial information as much as possible in spatial attention.
[0059] For the input feature X∈R C×H×W ,
[0060] When only channel attention is used:
[0061]
[0062] Among them, Wq, Wk, Wv are 1×1 convolutional layers, σ1, σ2 are reshape operators, and F SM is the softmax operator, and × is the matrix dot product operation The number of internal channels between Wv|Wq and Wz is C / 2;
[0063] The output of the channel branch only is: Among them, ch is the channel multiplication operator;
[0064] When only spatial attention is used:
[0065] A sp (X) = F SG [σ3(F SM (σ1(F GP (W q (X))))×σ2(W v (X)))]
[0066] Wq, Wv are standard 1×1 convolutions, θ2 is the intermediate parameter of these channel convolutions, σ1σ2σ3 are three reshape operators, F SM is the softmax operator, F GP (·) is the global pooling operator;
[0067]
[0068] The output of the space-only branch is: Among them, sp is the spatial multiplication operator;
[0069] The outputs of the above two branches are composed in a parallel layout:
[0070] PSA p (X) = Z ch +Z sp =A ch (X)⊙ ch X+A sp (X)⊙ sp X.
[0071] Example 2: The present invention is based on the PSA-Swin Transformer image classification algorithm and conducts the following experiments:
[0072] 1) Dataset and indicator evaluation
[0073] The Mini-Imagenet dataset was created by the Google DeepMind team in 2016 using a small portion of the Imagnet dataset. The dataset consists of 50,000 training images and 10,000 test images, which are evenly distributed in 100 categories, and the image size is not fixed.
[0074] The model settings are as follows: training was performed for 300 epochs using the AdamW optimizer, a batch size of 128, an initial learning rate of 0.001, and a weight decay of 0.05 using a cosine decay learning rate scheduler and a linear warmup of 20 epochs.
[0075] The evaluation indicator is accuracy (ACC).
[0076] 2) Performance comparison
[0077] First, the proposed model is compared with the current mainstream transformer model on the Mini-Imagenet dataset. The results show that the proposed model outperforms all reference baselines in terms of performance.
[0078] TABLE II
[0079] IMAGE CLASSIFICATION PERFORMANCE ON THE IMAGENET VALIDATION SET."PARAM" REFERS TO THE NUMBER OF PARAMETERS
[0080]
[0081] 3) Ablation analysis
[0082] In order to verify the functional significance of each component module in the architecture, the present invention conducted a detailed ablation study. The results are shown in the following table:
[0083] TABLE III
[0084] ABLATION STUDIES OF EPE AND PSA MODULES.EXPERIMENTS ARE CONDUCTED ONTHE MINI-IMAGENET DATASET
[0085] Model EPE PSA Batch Size ACC1(%)↑ PSA-Swin-T - - 128 82.48 PSA-Swin-T + - 128 84.82 PSA-Swin-T - + 128 82.98 PSA-Swin-T + + 128 85.14 PSA-Swin-S - - 128 83.07 PSA-Swin-S + - 128 85.63 PSA-Swin-S - + 128 83.74 PSA-Swin-S + + 128 86.03
[0086] An ablation study was performed at an image resolution of 224×224 to examine the effect of removing or adding the EPE and PSA modules. At a 224×224 input image resolution setting, the EPE module improved the accuracy by 2.34% compared to the SwinTransformer Tiny version without any modules. The PSA module improved the accuracy by 0.5% compared to the SwinTransformer Tiny version without any modules. Using both the EPE and PSA modules increased the accuracy by 2.66% to 85.14%.
[0087] In addition, under the 224×224 input image resolution setting, the EPE module improves the accuracy by 2.56% over the Swin Transformer Small version without any modules. The PSA module improves the accuracy by 0.67% compared to the Swin Transformer Small model without any modules. After using the EPE and PSA modules at the same time, the accuracy is improved by 2.96% to 86.03%. This shows that the EPE and PSA modules can significantly improve the accuracy of low-resolution images. This proves the importance of the two modules proposed in this invention, namely, efficient position encoding and the introduction of attention help improve the accuracy.
Claims
1. An image classification algorithm based on PSA-Swin Transformer, characterized in that: Includes steps: S1. First, the image is input into the Patch Partition module for block division, and then flattened in the channel direction; S2, construct feature maps of different sizes through four stages in sequence; In the first stage, the data first passes through a Linear Embedding layer and then passes through a PSA-Swin TransformerBlock layer. In the second, third, and fourth stages, the image is first downsampled through a Patch Merging layer and then through a PSA-Swin Transformer Block layer; The PSA-Swin Transformer Block layer includes the EPE module and the PSA module; S3, then connect a Layer Norm layer, a global pooling layer and a fully connected layer to get the final output.
2. The image classification algorithm based on PSA-Swin Transformer according to claim 1, characterized in that: In step S1, the image is an RGB three-channel image. Each patch has 4×4=16 pixels. Each pixel has three values of R, G, and B. After flattening, it is 16×3=48. After passing through the Patch Partition layer, the shape of the image changes from [H, W, 3] to [H / 4, W / 4, 48]. Then, through the Linear Embedding layer, a linear transformation is performed on the channel data of each pixel, that is, the original number of channels is changed from 48 to C, and the shape of the image is changed from [H / 4, W / 4, 48] to [H / 4, W / 4, C], where H represents the height of the image, W represents the width, and C represents the number of channels.
3. The image classification algorithm based on PSA-Swin Transformer according to claim 1, characterized in that: In step S2, when the feature map passes through the first Patch Merging layer, the Patch Merging layer divides each 2x2 adjacent pixel into a patch, and then splices the pixels at the same position in each patch together to obtain 4 times the number of the original feature map; then the feature map is spliced in the channel direction, and then passes through a layer of normalization layer, and finally passes through a fully connected layer to make a linear change in the depth direction of the feature map, so that the depth of the feature map will be changed from C at the input to 2C, and the resolution will be halved; after the second stage, the output resolution is H / 8×W / 8, and the number of channels becomes 2C; after the third stage, the output resolution is H / 16×W / 16, and the number of channels becomes 4C; The output resolution after the fourth stage is H / 32×W / 32, and the number of channels becomes 8C.
4. The image classification algorithm based on PSA-Swin Transformer according to claim 1, characterized in that: In step S2, EPE(X) = GSConv(X), where GSConv uses a shuffle operation to allow information exchange between groups; GSConv uses a grouped convolution architecture to first distribute input features to different groups, and then independently performs convolution operations within each group.
5. The image classification algorithm based on PSA-Swin Transformer according to claim 1, characterized in that: In step S2, the PSA module accurately achieves pixel-by-pixel regression by combining two key strategies: First, the polarization filtering strategy completely folds the input tensor along the corresponding dimension while maintaining the internal resolution when performing channel and spatial attention calculations; Second, the enhancement strategy is achieved by introducing nonlinearity in the attention mechanism.
6. The image classification algorithm based on PSA-Swin Transformer according to claim 5, characterized in that: In the PSA module, for the input feature X∈R C×H×W , When only channel attention is used: Among them, Wq, Wk, Wv are 1×1 convolutional layers, σ1, σ2 are reshape operators, and F SM is the softmax operator, and × is the matrix dot product operation The number of internal channels between Wv|Wq and Wz is C / 2; The output of the channel branch only is: Among them, ch is the channel multiplication operator; When only spatial attention is used: A sp (X)=F SG [σ3(Fs M (σ1(F GP (W q (X))))×σ2(W v (X)))] Wq, Wv are standard 1×1 convolutions, θ2 is the intermediate parameter of these channel convolutions, σ1σ2σ B There are three reshape operators, F SM is the softmax operator, F GP (·) is the global pooling operator; The output of the space-only branch is: Among them, sp is the spatial multiplication operator; The outputs of the above two branches are composed in a parallel layout: PSA p (X)=Z ch +Z sp =A ch (X)⊙ ch X+A sp (X)⊙ sp X。