关键词: ViT模型, 工业图像分类, CMT模型, 注意力, 特征融合

Abstract: The industrial images have some character, such as limited data acquisition, complex environments and variable lighting conditions, the classification accuracy of ViT model remains suboptimal. To address this issue, an industrial image classification algorithm is proposed based on the CMT model. Firstly, the Patch Embedding module is enhanced by incorporating affine transformations and sequential convolutional blocks, improving the  generalization capability of model on small datasets. Subsequently, the CMT Block is refined by introducing a parallel local feature extraction module, which enhances the model??s ability to capture local features. The multi-head self-attention mechanism is replaced with a token interaction attention mechanism to improve the model??s global feature representation. Deep convolution and channel attention are then integrated into the feedforward neural network, enabling the model to effectively capture neighboring features. Finally, a feature fusion module is proposed to integrate local and global features, enriching the feature representation and enhance classification performance on small datasets. Experimental results on the self-made filling bucket dataset, the public Car Parts dataset, and the Tiny ImageNet dataset demonstrate that the improved CMT model achieves classification accuracy improvements of 4.7, 6.9 and 5.2 persentage points for Top-1 Accuracy over the CMT model and 0.057, 0.071 and 0.048 for Macro F1 over the CMT model.

Key words: ViT model industrial image classification, CMT model, attention, feature fusion