从图像分类到目标检测：手把手教你用Hugging Face Transformers库玩转ViT和DETR
从图像分类到目标检测：手把手教你用Hugging Face Transformers库玩转ViT和DETR
视觉Transformer模型正在彻底改变计算机视觉领域。两年前，如果你告诉一个CV工程师所有图像任务都可以用Transformer解决，他可能会觉得你在开玩笑。但今天，从图像分类到目标检测，Transformer架构不仅证明了其有效性，而且在许多基准测试中超越了传统CNN方法。本文将带你快速上手两个最具代表性的视觉Transformer模型：ViT（Vision Transformer）和DETR（Detection Transformer），使用Hugging Face生态快速实现从零到生产的完整流程。
1. 环境准备与工具链搭建
在开始之前，我们需要配置一个高效的开发环境。推荐使用Python 3.8+和PyTorch 1.10+的组合，这是目前最稳定的搭配。
首先安装核心依赖库：
pip install torch torchvision torchaudio pip install transformers timm datasets对于GPU加速，建议安装对应CUDA版本的PyTorch。可以通过以下命令检查安装是否成功：
import torch print(torch.__version__) print(torch.cuda.is_available()) # 应该返回True常见问题排查：
- 如果遇到CUDA相关错误，尝试降低CUDA版本或重新安装PyTorch
timm库版本建议0.6.0以上，以获得完整的ViT支持- 内存不足时，可以添加
--no-cache-dir参数减少pip的内存占用
提示：使用虚拟环境（如conda或venv）可以避免依赖冲突。对于生产环境，建议将依赖固定到特定版本。
2. ViT实战：图像分类全流程
Vision Transformer（ViT）将图像分割为固定大小的patch，然后像处理NLP中的token一样处理这些patch。我们使用Hugging Face提供的预训练模型google/vit-base-patch16-224，这是一个在ImageNet-21k上预训练的基准模型。
2.1 快速推理示例
以下代码展示了如何使用ViT进行图像分类：
from transformers import ViTImageProcessor, ViTForImageClassification from PIL import Image import requests # 加载模型和处理器 processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224') model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224') # 准备输入图像 url = 'http://images.cocodataset.org/val2017/000000039769.jpg' image = Image.open(requests.get(url, stream=True).raw) # 预处理和预测 inputs = processor(images=image, return_tensors="pt") outputs = model(**inputs) logits = outputs.logits # 解析结果 predicted_class_idx = logits.argmax(-1).item() print("Predicted class:", model.config.id2label[predicted_class_idx])关键参数说明：
patch_size=16：图像分割的patch大小image_size=224：输入图像的预期尺寸num_attention_heads=12：Transformer的注意力头数
2.2 微调自定义数据集
要在自己的数据集上微调ViT，可以使用Hugging Face的TrainerAPI。以下是一个简化流程：
- 准备数据集（使用
datasets库） - 定义数据增强策略
- 配置训练参数
- 开始训练
from transformers import TrainingArguments, Trainer training_args = TrainingArguments( output_dir='./results', per_device_train_batch_size=16, evaluation_strategy="steps", num_train_epochs=3, save_steps=500, eval_steps=500, logging_dir='./logs', learning_rate=2e-5, ) trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, ) trainer.train()性能优化技巧：
- 使用混合精度训练（
fp16=True） - 梯度累积（
gradient_accumulation_steps）减少内存消耗 - 调整学习率调度器（默认为线性衰减）
3. DETR实战：端到端目标检测
DETR（Detection Transformer）是Facebook提出的端到端目标检测模型，消除了传统方法中的人工设计组件（如anchor）。我们使用facebook/detr-resnet-50预训练模型。
3.1 基础检测流程
from transformers import DetrImageProcessor, DetrForObjectDetection import torch processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") inputs = processor(images=image, return_tensors="pt") outputs = model(**inputs) # 将输出转换为COCO API格式 target_sizes = torch.tensor([image.size[::-1]]) results = processor.post_process_object_detection( outputs, target_sizes=target_sizes, threshold=0.9 )[0] for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): box = [round(i, 2) for i in box.tolist()] print( f"Detected {model.config.id2label[label.item()]} with confidence " f"{round(score.item(), 3)} at location {box}" )DETR独特优势：
- 无需NMS（非极大值抑制）后处理
- 固定数量的预测输出（默认100个）
- 全局上下文感知，减少重复检测
3.2 自定义数据训练
DETR的微调需要特别注意数据格式。建议使用COCO格式的标注，并通过以下方式加载：
from datasets import load_dataset dataset = load_dataset("cppe-5") # 示例数据集训练时需要自定义损失函数，因为DETR使用匈牙利匹配算法：
def collate_fn(batch): pixel_values = [item["pixel_values"] for item in batch] labels = [item["labels"] for item in batch] return {"pixel_values": torch.stack(pixel_values), "labels": labels} training_args = TrainingArguments( output_dir="./detr-finetuned", per_device_train_batch_size=4, num_train_epochs=10, save_steps=500, logging_steps=100, learning_rate=1e-5, remove_unused_columns=False, ) trainer = Trainer( model=model, args=training_args, data_collator=collate_fn, train_dataset=dataset["train"], eval_dataset=dataset["test"], ) trainer.train()训练注意事项：
- DETR训练收敛较慢，建议至少30个epoch
- 学习率不宜过大，1e-5是较好的起点
- 批量大小受限于GPU内存，可以小至2-4
4. ViT与DETR深度对比与应用选择
虽然ViT和DETR都基于Transformer架构，但它们在设计理念和应用场景上有显著差异。下面从多个维度进行对比：
| 特性 | ViT | DETR |
|---|---|---|
| 主要任务 | 图像分类 | 目标检测 |
| 输入处理 | 直接patch分割 | CNN backbone特征提取 |
| 位置编码 | 1D可学习或固定 | 2D正弦位置编码 |
| 输出结构 | 单一分类标签 | 固定数量预测框 |
| 训练复杂度 | 相对简单 | 需要匈牙利匹配 |
| 推理速度（224x224） | ~30ms/图像 | ~120ms/图像 |
| 典型应用场景 | 大规模图像分类 | 需要精确目标定位的任务 |
模型选择建议：
- 当只需要知道图像内容类别时，选择ViT
- 当需要定位图像中的多个对象时，选择DETR
- 对于实时性要求高的场景，可以考虑ViT的蒸馏版本（如
deit-tiny） - 对于小目标检测，DETR可能表现不如传统方法，可考虑后续改进模型如Deformable DETR
5. 高级技巧与性能优化
5.1 混合精度训练
同时使用FP16和FP32可以显著减少内存占用并加速训练：
training_args = TrainingArguments( fp16=True, ... )5.2 梯度检查点
对于大模型或有限显存的情况，可以启用梯度检查点：
model.gradient_checkpointing_enable()5.3 量化推理
使用8位量化可以减小模型体积并加速推理：
from transformers import quantization quantized_model = quantization.quantize_dynamic( model, {torch.nn.Linear}, dtype=torch.qint8 )5.4 模型蒸馏
从小型ViT开始可以大幅提升推理速度：
from transformers import DeiTForImageClassification small_model = DeiTForImageClassification.from_pretrained('facebook/deit-tiny-patch16-224')实际项目经验：
- 在部署到边缘设备时，量化能带来3-4倍的加速
- 蒸馏模型精度下降通常在1-3%以内，但体积缩小5-10倍
- 对于生产环境，建议构建模型服务而不是直接调用