Vision Transformer训练

### 如何训练 Vision Transformer (ViT) 模型

#### 数据准备
为了成功训练 ViT，数据集的选择至关重要。通常情况下，大规模的数据集如 ImageNet 或者更大的 JFT-300M 被用于预训练阶段[^1]。这些大型数据集有助于提高模型泛化能力和性能。

#### 配置环境与安装依赖库
在开始之前，确保已经配置好 Python 开发环境并安装必要的软件包。可以利用 PyTorch 和 Hugging Face 的 `transformers` 库来简化实现过程：

pip install torch torchvision transformers timm

#### 加载预定义架构
通过使用第三方库如 TIMM（PyTorch Image Models），可以直接加载预先构建好的 ViT 架构而无需手动编写每一层的具体细节:

import timm

model = timm.create_model('vit_base_patch16_224', pretrained=True)

这段代码创建了一个基于补丁大小为 16×16 并且输入图片尺寸为 224 × 224 像素的基础版 ViT 模型，并初始化权重来自于已有的预训练参数。

#### 定义损失函数和优化器
对于图像分类任务来说，交叉熵是一个常见的选择作为损失函数；而对于优化算法，则可以选择 AdamW 这样的方法因为它对正则化的处理效果较好:

import torch.optim as optim
from torch.nn import CrossEntropyLoss

criterion = CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=1e-5, weight_decay=0.01)

#### 准备数据管道
建立合适的数据增强策略以及高效的 DataLoader 是至关重要的一步。这不仅提高了最终模型的表现力还加快了收敛速度:

from torchvision.transforms import Compose, Resize, ToTensor, Normalize
from torch.utils.data import DataLoader
from datasets import load_dataset

def get_data_loader(batch_size=32):
    transform = Compose([
        Resize((224, 224)),
        ToTensor(),
        Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    dataset = load_dataset("image_folder", data_dir="path/to/dataset")
    train_set = dataset['train'].with_transform(transform)

    return DataLoader(train_set, batch_size=batch_size, shuffle=True)

此部分展示了如何设置一个简单的数据转换流程并将它们应用于来自本地文件夹中的自定义数据集上。

#### 实现训练循环
最后就是核心环节—实际执行训练的过程，在这里会迭代遍历整个 epoch 来更新网络参数直到达到满意的精度水平为止:

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)

data_loader = get_data_loader()

for epoch in range(num_epochs):  
    running_loss = 0.0
    
    for inputs, labels in data_loader:
        optimizer.zero_grad()
        
        outputs = model(inputs.to(device))
        loss = criterion(outputs.logits, labels.to(device))

        loss.backward()
        optimizer.step()

        running_loss += loss.item()
    
    print(f'[Epoch {epoch + 1}] Loss: {running_loss / len(data_loader)}')

上述脚本实现了基本的前向传播、反向传播机制并通过调整学习率等超参控制着整个训练进程。