前阵子有个朋友问我现在做图像分类是不是必须得上大模型传统的卷积网络是不是已经过时了。我说不一定但如果你手里数据量够、算力也凑合想追一下现在学术界和工业界都在聊的Transformer架构那ViTVision Transformer确实是一个绕不开的入口。而且说句实在话ViT本身的代码逻辑并不复杂比起调优一个深层的卷积网络它的结构反而更“直白”一些。这篇文章我就用实际跑通的代码带你从零搭一个ViT图像分类的完整流程包括环境准备、模型搭建、训练推理、再到最后的可视化输出全程不讲废话只讲实操。这篇文章适合谁看适合那些已经会用PyTorch跑过简单CNN模型、但对ViT还只停留在“听说过”阶段的开发者。当然如果你刚入门深度学习只要跟着步骤走也能把流程跑通只是中间某些细节你可能需要多花点时间消化。1. 方案选型与核心思路拆解1.1 为什么选ViT而不是继续用CNNViT的全称是Vision Transformer它把NLP领域大放异彩的Transformer结构搬到了图像领域。你不需要先精通Transformer所有细节只需要知道它的核心机制叫“自注意力”图像会被切成一个个小块每个小块类似NLP里的一个“词”图像分类就变成了“词序列”的分类。和CNN相比ViT最大的特点是没有卷积操作纯粹靠注意力机制来建立图像不同区域之间的关联。这样做的好处是理论上它可以捕捉到图像中距离很远但语义相关的特征比如一张图里左上角的车轮和右下角的车身。传统CNN靠堆叠卷积层扩大感受野也能做到但ViT在一开始就具备全局建模能力。代价也很明显就是它会比同等规模的CNN更吃数据和算力所以在中小数据集上如果直接从头训练ViT效果往往打不过ResNet。这也是为什么绝大多数ViT实战项目都是在ImageNet预训练权重的基础上做迁移学习。我选的方案就是“预训练ViT 微调”用一个具体的小型图像数据集来做分类演示。这样既能在有限算力下跑出不错的效果又能让你清晰看到ViT的训练流程和可视化效果。1.2 数据集与任务定义为了演示我们用某个开源的昆虫图像数据集——别纠结具体是哪个数据集因为整个流程是通用的换成你自己的数据也只是一行代码的事。数据集包含两类昆虫训练集和验证集分别放在对应目录下每个类一个文件夹。这个任务的核心就一句话给定一张昆虫图片模型要判断它属于哪个类别。听起来很简单但它涵盖了图像分类从数据加载、模型构建、训练验证到结果可视化的完整闭环。我会把每个环节的代码都拆开讲清楚。1.3 技术栈选型为什么是PyTorch HuggingFace现在做ViT绕不开的两个库是PyTorch和HuggingFace的transformers。PyTorch不用多说它的动态图机制和生态让自定义模型和调试都变得很舒服。HuggingFace的transformers库原本是NLP领域的标配后来也把视觉模型整合进来了使用统一的API加载ViT预训练权重非常方便。当然你也可以直接用timm库那里也有很多现成的ViT变体加载更轻量。我这次选择HuggingFace的版本是因为它的API风格很统一后面如果想切换到其他Transformer模型比如Swin Transformer代码改动量非常小。还有一个原因是HuggingFace的权重下载在国内有些网络环境下可能需要配置镜像我会在后面的常见问题里专门讲。2. 环境搭建与依赖准备2.1 Python环境与GPU检查第一步先准备一个干净的Python环境。我推荐用Anaconda管理避免不同项目之间的依赖冲突。新建环境的命令很简单conda create -n vit_demo python3.10 conda activate vit_demoPython版本我选3.10这个版本对PyTorch和transformers的兼容性都很稳定。接下来安装PyTorch。这里有个重要选择CPU版还是GPU版。如果你的电脑有NVIDIA显卡并且CUDA环境之前配过可以用GPU版训练速度能快一个数量级如果没有也可以先用CPU跑代码逻辑完全一样只是慢。# CPU版本 pip install torch torchvision # GPU版本以CUDA 12.1为例 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121装完后用一行代码验证GPU是否可用python -c import torch; print(torch.cuda.is_available())输出True说明GPU正常输出False就说明当前用的是CPU。我实测下来的经验是即使你的显卡只有4GB显存跑ViT-TinyViT的最小变体参数量约500万微调也完全够用。2.2 安装核心依赖库接下来安装transformers、datasets、以及一些辅助库pip install transformers datasets accelerate scikit-learn matplotlib pillow这里我额外装了一个accelerate它是HuggingFace的训练加速库后面用Trainer训练时会用到。datasets库是用来加载和处理数据集的虽然我们这次从本地文件夹读取图片但用datasets的ImageFolder功能可以省去不少写数据加载器的代码。装好之后推荐顺手装一下WandB或者TensorBoard二选一即可用来记录训练曲线。我这次用的是TensorBoard因为它是本地服务不需要注册账号。2.3 数据集预检查与目录结构动手写代码前先把数据集的目录结构确认好。一个规范的ImageFolder目录长这样data/ ├── train/ │ ├── class_a/ │ │ ├── img_001.jpg │ │ ├── img_002.jpg │ │ └── ... │ └── class_b/ │ ├── img_100.jpg │ └── ... └── val/ ├── class_a/ │ └── ... └── class_b/ └── ...每个类一个文件夹文件夹名就是类别标签。这个结构是PyTorch的ImageFolder和HuggingFace的load_dataset都能直接用的标准格式几乎不用写额外的数据解析代码。这里有一个我在真实项目中踩过的坑类文件夹的命名不要用中文或带空格有些库处理路径时容易出问题最好一律用英文小写加下划线。3. 核心代码实现与参数解析3.1 加载预训练模型与特征提取器ViT模型在送入Transformer之前图片要先被切成固定大小的patch并转成token序列。这个预处理过程需要和预训练时保持一致否则模型效果会大打折扣。HuggingFace把这部分封装在了AutoFeatureExtractor里直接调用就行。from transformers import AutoImageProcessor, AutoModelForImageClassification image_processor AutoImageProcessor.from_pretrained(google/vit-base-patch16-224) model AutoModelForImageClassification.from_pretrained( google/vit-base-patch16-224, num_labels2, ignore_mismatched_sizesTrue )这里我用的是vit-base-patch16-224意思是输入图片尺寸是224x224patch大小是16x16模型参数量约8600万。这个规格是目前最常用的基准配置效果和速度比较均衡。需要注意ignore_mismatched_sizesTrue这个参数。因为原始的vit-base模型是在ImageNet上做1000类分类而我们的任务只有2类输出层的shape不匹配加了这句话库会自动帮我们替换掉最后的分类头。3.2 数据增强与预处理流程数据增强是图像分类里非常重要的一个环节它相当于在训练时“免费”扩充数据集提升模型的泛化能力。对于ViT来说常用的增强手段和CNN类似但有一个细节ViT对输入图片的标准化要求更严格因为预训练时的像素分布是固定的。from transformers import AutoImageProcessor from torchvision import transforms image_processor AutoImageProcessor.from_pretrained(google/vit-base-patch16-224) train_transforms transforms.Compose([ transforms.RandomResizedCrop(image_processor.size[height]), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(meanimage_processor.image_mean, stdimage_processor.image_std) ]) val_transforms transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(meanimage_processor.image_mean, stdimage_processor.image_std) ])这里有个细节值得说一下训练时用RandomResizedCrop随机裁剪缩放相当于不断告诉模型“这个物体的主体不一定在图片正中央”验证时用Resize到256再CenterCrop到224是为了尽量保留图片中间区域的信息模拟真实推理场景。Normalize的均值和标准差我直接用了image_processor里自带的它们是ImageNet数据集的统计值不要自己随便改。3.3 数据集封装与加载我比较推荐把自定义数据集封装成一个PyTorch的Dataset类虽然代码多一点但后面调试和扩展都很方便。核心逻辑就是遍历目录下的所有图片路径和标签然后按索引取出图片并套用预处理。from torch.utils.data import Dataset from PIL import Image import os class InsectDataset(Dataset): def __init__(self, root_dir, transformNone): self.classes sorted(os.listdir(root_dir)) self.class_to_idx {cls: idx for idx, cls in enumerate(self.classes)} self.samples [] for cls in self.classes: cls_dir os.path.join(root_dir, cls) for fname in os.listdir(cls_dir): self.samples.append((os.path.join(cls_dir, fname), self.class_to_idx[cls])) self.transform transform def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label self.samples[idx] image Image.open(img_path).convert(RGB) if self.transform: image self.transform(image) return image, label这里有一个关键操作.convert(RGB)。有些图片可能是灰度图或者带透明通道的PNG如果不统一转成RGB后面送入模型时Tensor的shape会对不上报错很隐蔽。我也遇到过一种情况是图片本身损坏打不开这种情况下PIL会直接抛异常稳妥的做法是在__getitem__里加个try-except遇到坏图就跳过或者返回同目录下的其他图。3.4 训练配置与模型编译数据准备好之后接下来是训练的核心配置。这里我用HuggingFace的Trainer来做训练它把训练循环、梯度累积、日志记录都封装好了代码量能省一半。但为了让读者理解底层逻辑我先把关键的超参数列出来并逐一解释。from transformers import Trainer, TrainingArguments training_args TrainingArguments( output_dir./results, num_train_epochs10, per_device_train_batch_size16, per_device_eval_batch_size32, learning_rate5e-5, weight_decay0.01, eval_strategyepoch, save_strategyepoch, logging_dir./logs, logging_steps50, load_best_model_at_endTrue, metric_for_best_modelaccuracy, remove_unused_columnsFalse, )这些参数里面有四个是新手最容易忽略的learning_rate5e-5这个学习率比从头训练CNN时小得多因为预训练模型的权重已经挺好了我们只需要“微调”学习率太大容易把学好的特征破坏掉。说实话这个值我开始也调大过效果明显差一截。weight_decay0.01L2正则化的权重衰减系数帮助抑制过拟合。eval_strategyepoch每个epoch结束之后跑一次验证集这样能实时拿验证集指标判断当前模型好坏。load_best_model_at_endTrue训练结束后自动加载验证指标最好的一次模型权重而不是最后一次训练的权重。这个太有用了因为训练后期模型可能已经轻微过拟合看验证集指标它是在下降的最好的权重往往在中间某个位置。接着定义评估函数分类任务一般用准确率import numpy as np from sklearn.metrics import accuracy_score def compute_metrics(eval_pred): logits, labels eval_pred predictions np.argmax(logits, axis-1) return {accuracy: accuracy_score(labels, predictions)}然后创建Trainer对象trainer Trainer( modelmodel, argstraining_args, train_datasettrain_dataset, eval_datasetval_dataset, compute_metricscompute_metrics, )Trainer内部的默认数据整理逻辑需要把数据集返回的(image, label)转换成模型能接收的格式。这里我们传入了一个data_collator把图像像素堆叠成batchdef collate_fn(batch): images torch.stack([item[0] for item in batch]) labels torch.tensor([item[1] for item in batch]) return {pixel_values: images, labels: labels} trainer Trainer( modelmodel, argstraining_args, train_datasettrain_dataset, eval_datasetval_dataset, compute_metricscompute_metrics, data_collatorcollate_fn, )3.5 训练执行与显存管理执行训练只需要一行代码trainer.train()但在启动训练之前我还想啰嗦一句先跑一个batch的前向确认数据流没有报错。你可以简单地batch next(iter(train_loader)) output model(pixel_valuesbatch[0]) print(output.logits.shape)这一步能提前暴露很多问题比如图片尺寸不对、标签数量不匹配等等省得浪费一整轮训练时间。训练过程中你会看到类似这样的输出epoch 1/10, loss 0.8234, accuracy 0.6825 epoch 2/10, loss 0.5972, accuracy 0.7431 ...一般来说第3个epoch之后准确率会快速上升到第7个epoch左右开始趋于平缓10个epoch在小数据集上已经绰绰有余了。如果你发现loss一直不降大概率是学习率设置问题或者数据加载顺序不对优先检查数据增强那段代码。4. 可视化输出与结果分析4.1 损失曲线与准确率曲线训练完了光有终版数字还不够把曲线画出来才能直观判断有没有过拟合。Trainer训练过程的日志会存在training_args.logging_dir里我用TensorBoard直接可视化。打开一个新的终端tensorboard --logdir ./logs然后浏览器访问localhost:6006就能看到loss和accuracy曲线。按我的经验一个正常的训练曲线应该长这样训练集loss整体稳定下降。验证集loss在第5~7个epoch之间到达最低点之后开始反弹。这个反弹点就是过拟合开始的地方好在因为我们设了load_best_model_at_endTrue最终保存的权重是反弹之前那次。如果你不想用TensorBoard也可以用matplotlib把trainer.state.log_history里的数据手动画出来代码也很短适合丢在Notebook里直接看。4.2 混淆矩阵与分类错误分析准确率只能给一个宏观感受混淆矩阵才能告诉你模型在哪些类别之间容易混淆。这个分析在工业项目里非常重要因为它能帮你判断模型是不是真的学会了核心特征而只是靠背景记住标签。from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay import matplotlib.pyplot as plt # 得到验证集的预测结果 preds trainer.predict(val_dataset) y_pred np.argmax(preds.predictions, axis-1) y_true preds.label_ids cm confusion_matrix(y_true, y_pred) disp ConfusionMatrixDisplay(confusion_matrixcm, display_labelsval_dataset.classes) disp.plot(cmapBlues) plt.show()我分析过一次实际训练的结果发现模型把“甲虫A”误判为“甲虫B”的比例特别高。点开出错样本一看发现A类的大多数测试图片都带有绿色背景而B类的部分图片也有绿色背景模型很可能学了背景而不是昆虫本身。这就是为什么数据增强里加ColorJitter和RandomResizedCrop很重要它们能让模型少依赖颜色和位置信息。4.3 注意力图可视化ViT的“热力图”ViT比CNN多了一个优势就是它的注意力权重可以直接可视化让我们看到模型到底在关注图像的哪些区域。这个可视化是我平时最喜欢给朋友演示的部分因为效果非常直观。具体做法是前向传播的时候把模型最后一层Transformer的注意力权重拿出来然后按patch的位置映射回原图尺寸。import torch import matplotlib.pyplot as plt from PIL import Image import numpy as np def visualize_attention(image_path, model, image_processor): # 读取并预处理图片 image Image.open(image_path).convert(RGB) inputs image_processor(imagesimage, return_tensorspt) # 前向传播并获取注意力权重 with torch.no_grad(): outputs model(**inputs, output_attentionsTrue) attentions outputs.attentions[-1] # 取最后一层的注意力 # 取CLS token对所有patch的注意力权重形状为 (batch, heads, seq_len, seq_len) cls_attn attentions[0, :, 0, 1:].mean(dim0) # 对所有head取平均 num_patches int(cls_attn.shape[-1] ** 0.5) attn_map cls_attn.reshape(num_patches, num_patches).detach().numpy() # 上采样到原图尺寸 attn_map np.kron(attn_map, np.ones((14, 14))) # 叠加显示 fig, ax plt.subplots(1, 2, figsize(10, 5)) ax[0].imshow(image) ax[1].imshow(image) ax[1].imshow(attn_map, alpha0.6, cmapjet) plt.show()上面代码里有个关键点output_attentionsTrue这个参数它告诉模型在返回预测结果的同时把每一层Transformer的注意力矩阵也带出来。然后我取了最后一层、CLS token用来分类的那个特殊token对所有图像patch的注意力权重把所有注意力头平均一下reshape成patch网格的尺寸再上采样到原始分辨率。说实话第一次跑出这个热力图的时候我还是挺震撼的模型确实把大多数注意力集中到了昆虫身体的核心区域比如翅膀和触角而不是背景。这给调试模型提供了一个很有说服力的依据如果你的模型在错误分类时注意力集中在背景上那你就能判断是数据偏差问题而不是模型结构问题。5. 常见问题与排查技巧实录5.1 预训练权重下载慢或超时国内网络环境下HuggingFace的权重下载经常会卡住。解决方法是设置HF_ENDPOINT环境变量指向镜像站点。export HF_ENDPOINThttps://hf-mirror.com设置完之后重新运行脚本下载速度会明显提升。如果是公司防火墙比较严格也可以手动下载权重文件放到本地路径然后from_pretrained直接用本地路径。这里有个细节权重文件要放在一个单独文件夹里文件夹里还要有config.json文件否则库无法正确识别。5.2 显存不足怎么办如果你在训练或推理时报CUDA out of memory最简单的方案是减小per_device_train_batch_size从16降到8或4。梯度累积功能可以在不减小batch_size的情况下模拟大batchtraining_args TrainingArguments( per_device_train_batch_size8, gradient_accumulation_steps2, # 相当于batch_size16 )原理是每步只算8个样本的梯度攒够两步再更新一次参数效果接近直接用16的batch size。另一个常用技巧是使用半精度训练HuggingFace的Trainer里直接加一个参数training_args TrainingArguments( fp16True, )显存占用直接减半而且在我实测中半精度对最终精度的影响非常小。但如果你的显卡不支持半精度加速这个参数可能不生效需要检查一下。5.3 模型效果差连baseline都不如这种情况我在真实项目里碰过若干次排查顺序一般是这样的先看训练集loss有没有下降。如果训练集loss也不降说明模型没有学到东西优先检查代码逻辑看看数据标签是不是对上了。如果训练集loss降了但验证集准确率很低说明过拟合了增加数据增强、调小学习率、加大weight_decay。如果验证集也还行但是测试集差那就是数据分布偏了可能需要重新划分数据集或者做更贴近真实场景的数据增强。一个很容易被忽略的小问题图片输入尺寸对不对。ViT-base-f16是patch16224x224正好分成14x14个patch。如果你传进去一个非常规尺寸比如320x320虽然有些库会自动调整但如果不自动调整patch数量就会和模型不匹配直接报shape错误。5.4 Trainer训练时报错“columns mismatch”用Trainer时如果数据集返回的是一个dict里面包含多个字段比如image和labelTrainer默认会把所有字段都当输入传给模型容易报多余的columns错误。解决办法就是设置remove_unused_columnsFalse然后data_collator里只保留pixel_values和labels这两个字段。我在前面给的示例代码里已经包含了这个配置但如果你自己写训练循环就容易踩到这个坑。5.5 推理部署时如何提速训练完成后模型需要拿去用。最简方式是把模型转成torchscript或者ONNX推理速度能快不少。这里给一个ONNX导出的示例import torch dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, vit_insect.onnx, input_names[pixel_values], output_names[logits], dynamic_axes{pixel_values: {0: batch_size}}, )导出后可以用ONNXRuntime来跑推理。我实际对比过ONNX在CPU上的推理速度比PyTorch原生的eager模式快大约20%~40%具体看CPU型号和线程数。如果部署在移动端还可以考虑转成CoreML或TFLite。6. 项目扩展方向与进阶建议ViT这套流程跑通之后把它改成其他图像任务其实非常顺手。比如多标签分类只需要换掉损失函数把最后的sigmoid换成BCEWithLogitsLoss。目标检测接入DETR或者其他基于Transformer的检测头。图像检索直接取模型的CLS embedding作为图像特征向量算余弦相似度。有两个方向我觉得是最值得试的。第一个就是把ViT换成分层结构模型比如Swin Transformer它兼顾了CNN的分层思想和Transformer的全局建模在ImageNet上的表现普遍比ViT原版更好。代码改动极小只需要把from_pretrained的模型名换一下image_processor不变。第二个是尝试在更大的数据集上做自监督预训练比如用MAEMasked Autoencoder的方式先预训练一个模型再在自己的小数据集上微调。这种做法在数据量少的情况下效果比直接用ImageNet预训练权重还要好。我自己在实践中体会最深的一点是ViT的源码看似简单但真正有价值的其实是对数据流和注意力机制的理解。你把它拆开看清楚之后会发现很多AI技巧在这个框架里都可以举一反三。希望你也能在跑通这个流程之后顺手去改改那些超参数看看注意力热力图的变化那样你才算真的把这套东西用起来了。
