STB-GraCapsNet: A Novel Capsule Network Structure with Swin Transformer Block
摘要
Capsule network is a new type of neural network encoding features into capsules and constructing the part-whole relationships, which demonstrated good performance in image classification. However, it has some issues such as weak feature extraction ability and large number of training parameters, leading to performance decline when processing images with complex background. To address these issues, this article proposes a novel capsule network structure called STB-GraCapsNet, with key improvements as follows: 1) It uses Swin Transformer Block with the shift window attention mechanism for data preprocessing to obtain richer features; 2) the features are then processed through Linear Embedding and passed into the Primary Capsule Layer to generate high-dimension feature capsules, which reduces primary capsules number and improves feature extraction ability of the network. Ablation and comparative experiments based on multiple datasets show that STB-GraCapsNet performs better on complex datasets: 77.8 \(\%\) and 36.2 \(\%\) reduction in parameters number of capsule structure, and 12.07 \(\%\) and 5.22 \(\%\) improvement in model accuracy, as compared to CapsNet and GraCapsNet, respectively.