ViT-H / CIFAR-10

使用 Vision Transformer Huge（ViT-H/14）在 CIFAR-10 上进行图像分类训练。

模型规格

参数	值
架构	ViT-H/14
输入分辨率	336 × 336（双三次插值）
图块大小	14 × 14
图块数量	576（24×24）
嵌入维度	1280
Transformer 层数	32
注意力头数	16
MLP 维度	5120
参数量	~632M
分类数	10（CIFAR-10）

训练设置

参数	值
数据集	CIFAR-10
图像尺寸	32×32 → 336×336（双三次插值）
Batch size	32
优化器	AdamW (lr=3e-4, weight_decay=0.05)
学习率调度	CosineAnnealingLR
损失函数	CrossEntropyLoss (label_smoothing=0.1)
梯度裁剪	max_norm=1.0
数据增强	RandomHorizontalFlip + RandAugment
Early Stopping	patience=5（test acc 连续 5 epoch 不提升即停止）
最大 Epoch	100

文件结构

ViT-H/
├── model.py          # ViT-H 模型定义
├── train.py          # 训练脚本（含早停、曲线保存）
├── README.md
├── data/             # CIFAR-10 数据集（自动下载）
└── outputs/
    ├── vit_h_best.pt        # 最优测试精度的 checkpoint
    ├── vit_h_final.pt       # 最终 epoch 的 checkpoint（含完整训练历史）
    └── training_curves.png  # 训练/测试 Loss & Accuracy 曲线

快速开始

安装依赖

pip install torch torchvision matplotlib

开始训练

python train.py

CIFAR-10 数据集会自动下载至 ./data/，训练结果保存在 ./outputs/。

加载已保存的模型

import torch
from model import ViTH

model = ViTH(num_classes=10)

# 加载最优 checkpoint
ckpt = torch.load("outputs/vit_h_best.pt", map_location="cpu")
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

print(f"Best test acc: {ckpt['test_acc']:.4f}  (epoch {ckpt['epoch']})")

输出说明

• vit_h_best.pt：测试集精度最高时保存的权重
• vit_h_final.pt：训练结束时保存，额外包含完整的 loss / acc 历史列表
• training_curves.png：每个 epoch 更新一次，包含 Loss 和 Accuracy 两张子图

Early Stopping 逻辑

测试集（直接作为验证集）accuracy 连续 5 个 epoch 无提升时自动终止训练，同时保存最优权重。

patience = 0
if test_acc > best_acc:
    保存最优模型, patience = 0
else:
    patience += 1
    if patience >= 5: 停止

注意事项

• ViT-H 参数量约 632M，建议在 GPU 上训练（CUDA / MPS 均支持）
• CPU 训练速度极慢，仅用于验证代码正确性
• 336×336 输入会带来较大显存开销，如果显存不足可减小 BATCH_SIZE