Vision Transformer (ViT) for CIFAR-10

这是一个使用 PyTorch 实现的 Vision Transformer (ViT) 模型，用于 CIFAR-10 图像分类任务。

项目结构

ViT/
├── ViT_backbone.py          # ViT 模型核心实现
├── train_vit_cifar10.py     # CIFAR-10 完整训练脚本
├── train_vit_simple.py      # 简化版训练脚本
├── training_curves.png      # 完整版训练曲线
├── training_curves_simple.png # 简化版训练曲线
├── __init__.py             # Python 包初始化文件
└── README.md               # 项目说明文档

模型特点

• Vision Transformer 架构: 基于 Transformer 的图像分类模型
• 多头自注意力机制: 实现了完整的多头注意力层
• 位置编码: 为图像块添加位置信息
• 分类头: 用于 CIFAR-10 的 10 类分类

依赖环境

torch>=1.9.0
torchvision>=0.10.0
matplotlib
numpy
tqdm
einops

安装依赖

pip install torch torchvision matplotlib numpy tqdm einops

使用方法

训练模型

1. 完整版训练 (推荐):

python train_vit_cifar10.py

2. 简化版训练:

python train_vit_simple.py

模型参数

• 图像尺寸: 224x224 (从 CIFAR-10 的 32x32 上采样)
• 块大小: 16x16
• 嵌入维度: 768
• 注意力头数: 12
• Transformer 层数: 12
• MLP 隐藏层维度: 3072

训练结果

项目包含了训练过程的可视化曲线图:

• training_curves.png: 完整版训练的损失和准确率曲线
• training_curves_simple.png: 简化版训练的损失和准确率曲线

模型架构详解

1. 图像预处理 (pre_process)

• 将图像分割成固定大小的块 (patches)
• 线性嵌入到指定维度
• 添加位置编码和 CLS token

2. Transformer 编码器 (Transformer_Encoder)

• 多头自注意力机制
• 前馈神经网络 (MLP)
• 层归一化和残差连接

3. 分类头

• 使用 CLS token 的输出进行分类
• 最终输出 10 个类别的概率分布

注意事项

1. 模型文件: 由于 GitHub 文件大小限制，预训练的 .pth 模型文件未包含在仓库中
2. 数据集: CIFAR-10 数据集会在首次运行时自动下载
3. 计算资源: 建议使用 GPU 进行训练以获得更好的性能

参考文献

• Dosovitskiy, A., et al. "An image is worth 16x16 words: Transformers for image recognition at scale." ICLR 2021.

作者

httkxy

许可证

MIT License