【亲测免费】 Vision Transformer (ViT) 使用教程 - PyTorch 实现

1. 项目介绍

Vision Transformer（ViT）是由Alexey Dosovitskiy等人提出的一种利用Transformer架构进行图像识别的方法。在论文《An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》中，他们展示了Transformers如何在计算机视觉任务上取得优秀效果，即使在相对较小的数据集上。本项目是lucidrains/vit-pytorch的实现，提供了一个简洁易用的PyTorch接口来复现ViT。

2. 项目快速启动

首先确保已经安装了PyTorch。若未安装，可以使用以下命令：

pip install torch torchvision

接下来，克隆项目到本地并安装依赖：

git clone https://github.com/lucidrains/vit-pytorch.git
cd vit-pytorch
pip install -r requirements.txt

以下是一个简单的ViT模型实例化和前向传播的例子：

import torch
from vit_pytorch import ViT

# 初始化一个ViT模型
model = ViT(num_classes=1000, image_size=(256, 128), patch_size=(32, 16),
            dim=1024, depth=6, heads=16, mlp_dim=2048, dropout=0.1, emb_dropout=0.1)

# 假设我们有一个随机生成的输入图像
img = torch.randn(1, 3, 256, 128)

# 运行模型进行预测
preds = model(img)

3. 应用案例和最佳实践

示例：自定义数据集的预处理

如果要使用自定义数据集，你需要将图片转化为与ViT兼容的格式。ViT将输入图像切分成固定大小的patch，并将这些patch展平作为Transformer的输入序列。

import torchvision.transforms as transforms

def preprocess_image(image_path):
    transform = transforms.Compose([
        transforms.Resize((256, 256)),  # 将图片调整为固定大小
        transforms.CenterCrop(256),   # 截取中心区域以保持纵横比
        transforms.ToTensor(),         # 转换为张量
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 对照ImageNet预训练标准归一化
    ])
    
    return transform(Image.open(image_path)).unsqueeze(0)


# 加载一张图片并预处理
img_tensor = preprocess_image('path_to_your_image.jpg')
preds = model(img_tensor)

最佳实践：训练与优化

使用AdamW优化器和适当的学习率调度策略可以帮助改善模型的训练过程：

from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR

optimizer = AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
scheduler = LambdaLR(optimizer, lambda epoch: max(0.1, 1 - (epoch + 1) / epochs))

for epoch in range(epochs):
    # 训练循环...
    scheduler.step()

4. 典型生态项目

• timm: 由Ross Wightman维护的一个广泛的PyTorch图像模型集合，包含了ViT的多种变体。https://github.com/rwightman/pytorch-image-models
• DeiT (Data-efficient Image Transformers): Facebook AI Research的团队提出了一种数据效率更高的ViT变体。https://github.com/facebookresearch/deit

以上就是ViT-PyTorch的基本使用和一些最佳实践，祝你在使用过程中取得好成果！