企业官网建设流程全解析

前阵子有个朋友问我，现在做图像分类是不是必须得上大模型，传统的卷积网络是不是已经过时了。我说不一定，但如果你手里数据量够、算力也凑合，想追一下现在学术界和工业界都在聊的Transformer架构，那ViT（Vision Transformer）确实是一个绕不开的入口。而且说句实在话，ViT本身的代码逻辑并不复杂，比起调优一个深层的卷积网络，它的结构反而更“直白”一些。这篇文章我就用实际跑通的代码，带你从零搭一个ViT图像分类的完整流程，包括环境准备、模型搭建、训练推理、再到最后的可视化输出，全程不讲废话，只讲实操。

这篇文章适合谁看？适合那些已经会用PyTorch跑过简单CNN模型、但对ViT还只停留在“听说过”阶段的开发者。当然，如果你刚入门深度学习，只要跟着步骤走，也能把流程跑通，只是中间某些细节你可能需要多花点时间消化。

1. 方案选型与核心思路拆解

1.1 为什么选ViT而不是继续用CNN

ViT的全称是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 python=3.10 conda activate vit_demo

Python版本我选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-Tiny（ViT的最小变体，参数量约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_labels=2, ignore_mismatched_sizes=True )

这里我用的是vit-base-patch16-224，意思是输入图片尺寸是224x224，patch大小是16x16，模型参数量约8600万。这个规格是目前最常用的基准配置，效果和速度比较均衡。

需要注意ignore_mismatched_sizes=True这个参数。因为原始的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(brightness=0.2, contrast=0.2), transforms.ToTensor(), transforms.Normalize(mean=image_processor.image_mean, std=image_processor.image_std) ]) val_transforms = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=image_processor.image_mean, std=image_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, transform=None): 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_epochs=10, per_device_train_batch_size=16, per_device_eval_batch_size=32, learning_rate=5e-5, weight_decay=0.01, eval_strategy="epoch", save_strategy="epoch", logging_dir="./logs", logging_steps=50, load_best_model_at_end=True, metric_for_best_model="accuracy", remove_unused_columns=False, )

这些参数里面，有四个是新手最容易忽略的：

• learning_rate=5e-5：这个学习率比从头训练CNN时小得多，因为预训练模型的权重已经挺好了，我们只需要“微调”，学习率太大容易把学好的特征破坏掉。说实话这个值我开始也调大过，效果明显差一截。
• weight_decay=0.01：L2正则化的权重衰减系数，帮助抑制过拟合。
• eval_strategy="epoch"：每个epoch结束之后跑一次验证集，这样能实时拿验证集指标判断当前模型好坏。
• load_best_model_at_end=True：训练结束后，自动加载验证指标最好的一次模型权重，而不是最后一次训练的权重。这个太有用了，因为训练后期模型可能已经轻微过拟合，看验证集指标它是在下降的，最好的权重往往在中间某个位置。

接着定义评估函数，分类任务一般用准确率：

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( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=val_dataset, compute_metrics=compute_metrics, )

Trainer内部的默认数据整理逻辑，需要把数据集返回的(image, label)转换成模型能接收的格式。这里我们传入了一个data_collator，把图像像素堆叠成batch：

def 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( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=val_dataset, compute_metrics=compute_metrics, data_collator=collate_fn, )

3.5 训练执行与显存管理

执行训练只需要一行代码：

trainer.train()

但在启动训练之前，我还想啰嗦一句：先跑一个batch的前向，确认数据流没有报错。你可以简单地：

batch = next(iter(train_loader)) output = model(pixel_values=batch[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_end=True，最终保存的权重是反弹之前那次。

如果你不想用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_matrix=cm, display_labels=val_dataset.classes) disp.plot(cmap="Blues") 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(images=image, return_tensors="pt") # 前向传播并获取注意力权重 with torch.no_grad(): outputs = model(**inputs, output_attentions=True) attentions = outputs.attentions[-1] # 取最后一层的注意力 # 取CLS token对所有patch的注意力权重，形状为 (batch, heads, seq_len, seq_len) cls_attn = attentions[0, :, 0, 1:].mean(dim=0) # 对所有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, alpha=0.6, cmap="jet") plt.show()

上面代码里有个关键点：output_attentions=True这个参数，它告诉模型在返回预测结果的同时，把每一层Transformer的注意力矩阵也带出来。然后我取了最后一层、CLS token（用来分类的那个特殊token）对所有图像patch的注意力权重，把所有注意力头平均一下，reshape成patch网格的尺寸，再上采样到原始分辨率。

说实话，第一次跑出这个热力图的时候我还是挺震撼的，模型确实把大多数注意力集中到了昆虫身体的核心区域，比如翅膀和触角，而不是背景。这给调试模型提供了一个很有说服力的依据：如果你的模型在错误分类时，注意力集中在背景上，那你就能判断是数据偏差问题，而不是模型结构问题。

5. 常见问题与排查技巧实录

5.1 预训练权重下载慢或超时

国内网络环境下，HuggingFace的权重下载经常会卡住。解决方法是设置HF_ENDPOINT环境变量，指向镜像站点。

export HF_ENDPOINT=https://hf-mirror.com

设置完之后重新运行脚本，下载速度会明显提升。如果是公司防火墙比较严格，也可以手动下载权重文件放到本地路径，然后from_pretrained直接用本地路径。这里有个细节，权重文件要放在一个单独文件夹里，文件夹里还要有config.json文件，否则库无法正确识别。

5.2 显存不足怎么办

如果你在训练或推理时报CUDA out of memory，最简单的方案是减小per_device_train_batch_size，从16降到8或4。梯度累积功能可以在不减小batch_size的情况下模拟大batch：

training_args = TrainingArguments( per_device_train_batch_size=8, gradient_accumulation_steps=2, # 相当于batch_size=16 )

原理是每步只算8个样本的梯度，攒够两步再更新一次参数，效果接近直接用16的batch size。另一个常用技巧是使用半精度训练，HuggingFace的Trainer里直接加一个参数：

training_args = TrainingArguments( fp16=True, )

显存占用直接减半，而且在我实测中，半精度对最终精度的影响非常小。但如果你的显卡不支持半精度加速，这个参数可能不生效，需要检查一下。

5.3 模型效果差，连baseline都不如

这种情况我在真实项目里碰过若干次，排查顺序一般是这样的：

• 先看训练集loss有没有下降。如果训练集loss也不降，说明模型没有学到东西，优先检查代码逻辑，看看数据标签是不是对上了。
• 如果训练集loss降了，但验证集准确率很低，说明过拟合了，增加数据增强、调小学习率、加大weight_decay。
• 如果验证集也还行，但是测试集差，那就是数据分布偏了，可能需要重新划分数据集，或者做更贴近真实场景的数据增强。

一个很容易被忽略的小问题：图片输入尺寸对不对。ViT-base-f16是patch16，224x224正好分成14x14个patch。如果你传进去一个非常规尺寸（比如320x320），虽然有些库会自动调整，但如果不自动调整，patch数量就会和模型不匹配，直接报shape错误。

5.4 Trainer训练时报错“columns mismatch”

用Trainer时，如果数据集返回的是一个dict，里面包含多个字段（比如image和label），Trainer默认会把所有字段都当输入传给模型，容易报多余的columns错误。解决办法就是设置remove_unused_columns=False，然后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不变。第二个是尝试在更大的数据集上做自监督预训练，比如用MAE（Masked Autoencoder）的方式先预训练一个模型，再在自己的小数据集上微调。这种做法在数据量少的情况下，效果比直接用ImageNet预训练权重还要好。

我自己在实践中体会最深的一点是：ViT的源码看似简单，但真正有价值的其实是对数据流和注意力机制的理解。你把它拆开看清楚之后，会发现很多AI技巧在这个框架里都可以举一反三。希望你也能在跑通这个流程之后，顺手去改改那些超参数，看看注意力热力图的变化，那样你才算真的把这套东西用起来了。