简介这是一套FastVIT图像分类实战资源面向希望用Transformer架构高效完成图像分类的深度学习初学者与开发者。内容围绕完整项目流程展开覆盖数据准备、模型训练、权重导出与测试评估等关键环节特别适合结合代码理解ViT优化版本在真实数据集上的落地方法。压缩包共2000个文件大小约764.79MB核心为10个Python脚本、1979张PNG样本图片另含PyTorch模型权重文件pt/pth、配置文件json及标签说明txt结构清晰便于按步骤学习。目前已有648人学习具备一定参考价值。通过对数据预处理、数据增强、FastVIT模型搭建、交叉熵训练及验证测试等模块的动手实践读者可以直观掌握图像分类任务从数据到部署的完整链路同时借助附带的结果文件检查预测输出加深对Transformer在图像识别中应用的理解。1. FastVIT图像分类这个快字是怎么省出来的FastVIT不是把ViT的Transformer模块删掉几个而是在保持Stem → Body → Head整体结构不变的前提下把标准的softmax注意力替换成线性注意力配合重参数化技巧降低推理成本。作为图像分类算法里相对较新的Transformer模型它在显存和速度上的优势很明显。我在一块只有6G显存的卡上跑ViT-Base做图像分类batch size只能撑到8换成FastVIT同配置能跑到32验证集准确率也没明显掉点。这篇文章拆的是一套完整的FastVIT图像分类流程makedata.py做数据准备、train.py训练、export_model.py导出、test.py测试外加class.json和result.json两个关键产物。适合想在资源有限环境里跑图像分类、又想深入理解Transformer落地细节的从业者。下面按脚本执行顺序讲重点放在参数怎么调、哪里容易翻车。2. 数据准备makedata.py怎么把图片变成可训练的数据集2.1 图像分类的数据集目录约定与class.json来源跑makedata.py之前先明确FastVIT的输入格式。图像分类任务里最通用的做法是ImageFolder约定train和val两个根目录下每个类别一个子目录子目录名就是类别标签。项目里给的示例图片数量不多但结构是一致的。手动维护的目录大致长这样data/ └── images/ ├── cat/ │ ├── 5e4d1ee0d.png │ └── 77291b3ad.png └── dog/ ├── 0367e0199.png └── 5a8b75712.png类别名建议只用英文字母和下划线不要用中文也不要带空格。中文类别名在Windows下文件路径编码容易出问题带空格的目录名在后续脚本解析时也可能造成意外。class.json就是从这个目录结构生成的把类别名按固定顺序映射成整数索引import json import os src_dir data/images classes sorted([d for d in os.listdir(src_dir) if os.path.isdir(os.path.join(src_dir, d))]) class_map {str(i): cls for i, cls in enumerate(classes)} with open(class.json, w, encodingutf-8) as f: json.dump(class_map, f, ensure_asciiFalse, indent2)这里我习惯对类别名做一次sorted保证顺序稳定。如果不排序每次运行都可能因为文件系统返回顺序不一致而生成不同的映射训练和推理时class.json的对应关系就乱了。class.json生成后建议人工打开扫一眼确认索引和类别名没有错位这个文件在后面test.py推理时直接决定预测结果翻译成什么标签错了就是全错。2.2 makedata.py的核心逻辑与参数说明运行makedata.py是创建训练集和验证集的第一步。脚本内部大致做三件事遍历原始图片目录、做尺寸归一化、按比例切分训练集和验证集。下面是我按常见做法补全的核心逻辑import os import random from PIL import Image def process(src_dir, train_dir, val_dir, val_ratio0.2, target_size(224, 224)): random.seed(42) for cls in os.listdir(src_dir): cls_path os.path.join(src_dir, cls) if not os.path.isdir(cls_path): continue files [f for f in os.listdir(cls_path) if f.lower().endswith((.png, .jpg, .jpeg))] random.shuffle(files) val_cnt int(len(files) * val_ratio) for i, f in enumerate(files): img Image.open(os.path.join(cls_path, f)).convert(RGB) img img.resize(target_size, Image.BILINEAR) out_dir os.path.join(val_dir if i val_cnt else train_dir, cls) os.makedirs(out_dir, exist_okTrue) img.save(os.path.join(out_dir, f))参数说明几点。val_ratio0.2是验证集比例数据量越小这个值应该越大几百张图时我倾向用0.3保证验证集里每个类都有足够样本。Image.BILINEAR是双线性插值速度和画质居中对图像质量要求高就换Image.BICUBIC代价是处理时间变长。convert(RGB)很关键灰度图和RGBA图如果不统一转成三通道后面训练时会报张量维度不匹配的错误。这个脚本默认不做数据增强增强放到train.py加载数据时实时做。把数据准备和增强分开的好处是makedata.py只负责切分和归一化训练时的增强策略可以随时调整不用重新跑一遍数据预处理。切分时按类别独立shuffle避免某一类图片特别少时验证集里缺了这个类。2.3 数据增强训练集和验证集为什么不能共用一套变换图像分类算法的数据增强手段不少FastVIT常用的训练时增强是随机裁剪加翻转验证时则只用固定缩放和中心裁剪。原因很简单训练时用随机性让模型见过更多变化验证时则要尽量贴近真实分布保证结果稳定可对比。from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.8, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])RandomResizedCrop的scale0.8到1.0表示裁剪区域占原图面积的比例。数据集小时可以放宽到0.5增强强度加大但太小会裁掉太多上下文模型反而容易学偏。ColorJitter的亮度、对比度各0.2是温和档深色或曝光不匀的数据集可以再降一点。Normalize必须用ImageNet的均值方差因为后面的FastVIT加载的是ImageNet预训练权重输入分布要和预训练时对齐。如果从头训练最好按自己的数据集重新统计mean和std否则开头几个epoch的loss会飘得比较厉害。数据准备好之后用torchvision的ImageFolder加载再用DataLoader组织成batchfrom torch.utils.data import DataLoader from torchvision.datasets import ImageFolder train_dataset ImageFolder(data/train, transformtrain_transform) val_dataset ImageFolder(data/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue)num_workers4是CPU多进程加载Linux下没问题Windows下偶尔会报EOFError遇到就把num_workers改成0或2。shuffle在训练集必须开验证集不能开否则验证集顺序每次不同会影响判断。batch_size要看显存定FastVIT t8在6G卡上开32基本安全换更大模型就降到16或8。3. 训练核心train.py里的FastVIT配置与循环细节3.1 FastVIT模型变体怎么选FastVIT官方实现里按计算量从低到高有fastvit_t8、fastvit_t12、fastvit_s12、fastvit_m16四个主要变体。它们共享同一套架构设计区别主要在Transformer块的宽度和深度。换变体只需要改一个字符串不需要动训练脚本里的其他逻辑。模型变体计算量适用场景fastvit_t8低快速验证流程、边缘设备fastvit_t12中低中等规模数据集、追求训练速度fastvit_s12中通用图像分类任务fastvit_m16高精度优先、GPU资源充足选型时不要盲目上最大的模型。我在小数据集上用m16做过对比因为过拟合严重最终准确率反而不如t12。模型容量要跟数据量匹配每个类别样本少于500张时优先考虑t8或t12数据量上万再考虑s12以上。训练中初始化模型的写法from fastvit import create_model model create_model( fastvit_t8, pretrainedTrue, num_classes2, # 必须和class.json的类别数一致 dropout0.1 )pretrainedTrue会加载ImageNet预训练权重。这个参数对收敛速度影响非常大小数据集从随机初始化开始训经常几百个epoch都难收敛加载预训练后几十个epoch就能到可用精度。dropout0.1是分类头前的随机失活比例数据集越小可以适当调大到0.3。3.2 损失函数、优化器与学习率调度FastVIT做图像分类常规组合是交叉熵损失、AdamW优化器、余弦退火学习率。这三件套配合预训练权重是Transformer类视觉模型比较不容易翻车的配置。import torch from torch import nn from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR criterion nn.CrossEntropyLoss() optimizer AdamW(model.parameters(), lr3e-4, weight_decay0.05) scheduler CosineAnnealingLR(optimizer, T_max30)学习率从3e-4起步比较稳数据量小降到1e-4数据量大可以尝试5e-4。weight_decay0.05是ViT训练范式里的经验值作用对象是权重本身bias和LayerNorm参数一般不参与。T_max30表示余弦退火半周期长度通常和总epoch数一致如果总epoch是60T_max设30的话学习率会在30个epoch时降到最低点然后回升这种周期策略在很多任务上效果更好。类别不平衡时给交叉熵加权重class_weights torch.tensor([1.0, 5.0]).to(device) criterion nn.CrossEntropyLoss(weightclass_weights)权重按类别样本数反比设定少数类权重放大模型不会只预测多数类。但权重别拉到10以上我见过权重过大导致loss振荡、验证集崩掉的情况一般单边权重不超过5。3.3 训练循环、验证循环与最优模型保存训练循环的骨架是前向传播 - 算loss - 反传 - 更新参数每个epoch结束跑一次验证。模型需要先移到device上否则默认在CPU上跑6万张图训到天亮都训不完。device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) epochs 30 best_acc 0.0 for epoch in range(epochs): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() # 验证 model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, preds torch.max(outputs, dim1) correct (preds labels).sum().item() total labels.size(0) acc correct / total scheduler.step() print(fEpoch {epoch1}/{epochs} | Loss: {running_loss/len(train_loader):.4f} | Val Acc: {100*acc:.2f}%) if acc best_acc: best_acc acc torch.save(model.state_dict(), best_model.pth)这里有两处新手最容易忽略。第一optimizer.zero_grad()必须在每次step前调用PyTorch的梯度会默认累加忘了清会造成loss曲线周期性跳变。第二model.eval()和torch.no_grad()要一起用。eval()切换BN和Dropout的工作模式no_grad()关闭自动求导两个都不加的话验证阶段显存占用几乎翻倍小显存卡直接OOM。判断是否收敛有个实用办法train_loss持续下降但val_acc不再上升说明过拟合了提前停在best_acc对应的epoch即可train_loss和val_acc都不动优先怀疑学习率太低或模型容量不足再去动数据增强。3.4 训练日志写入result.json训练过程中把每个epoch的loss和acc记录下来最后统一写进result.json。这样跑完不用翻控制台输出一条命令就能看到训练曲线import json results [] # 在每个epoch的验证结束后追加 results.append({ epoch: epoch 1, train_loss: round(running_loss / len(train_loader), 4), val_acc: round(acc, 4) }) with open(result.json, w, encodingutf-8) as f: json.dump(results, f, indent2, ensure_asciiFalse)result.json的作用不止是记录调参对比时特别有用。train_loss一路降、val_acc不升反降基本就是过拟合需要加强数据增强或提前停train_loss和val_acc都不动优先怀疑学习率太低或模型容量不足。读这个文件不要用记事本翻直接用pandas读成表格画loss曲线更直观。4. FastVIT避坑记录五个实际跑出来的问题这一节列的是我在FastVIT图像分类实战中真实遇到的五个问题按出现频率从高到低排。前两个最隐蔽因为它们不会直接报错只是模型效果差后三个报错信息比较明显但解决方向容易走偏。排查时建议先看错误信息再看数据最后才怀疑模型结构。4.1 显存OOM验证阶段忘了关计算图现象训练到中途报CUDA out of memory但nvidia-smi一看显存用量并不高靠前的epoch都没问题。原因验证循环里只写了model.eval()没有包torch.no_grad()验证时前向传播的计算图被完整保留在显存里。另一个常见原因是验证batch size设得比训练时还大验证阶段成了压垮显存的最后一根稻草。解决验证循环一律写成with torch.no_grad():包裹同时把验证batch_size调成训练时的一半甚至直接降成1也可以。还有一个排查手段用torch.cuda.max_memory_allocated()看一眼峰值显存确认是哪个阶段的占用涨上去了。如果调小batch size还是OOM用torch.cuda.empty_cache()释放碎片显存再不行就nvidia-smi看看是不是其他进程占着显存。4.2 resize直接拉伸把图像弄变形分类精度上不去现象测试集准确率稳定在低水平训练loss看起来也降不动回看预处理后的图片发现物体被压扁或拉长。原因makedata.py里用了img.resize((224, 224))强行把不规则的图片拉伸到固定尺寸宽高比被破坏猫变成了瘦长条车变成了扁盒子特征分布被扭曲。解决先等比缩放到短边256再CenterCrop裁224×224。代码上就是val_transform里那种Resize(256) CenterCrop(224)组合训练集用RandomResizedCrop自带等比裁剪。注意推理阶段和测试阶段必须用同一套预处理变换不能训练用随机裁剪、推理用整图resize模型看到的分布完全不同精度会掉一大截。4.3 pretrained权重下载失败卡在网络而不是代码现象运行train.py时报下载错误或RuntimeError提示pretrained weights加载失败整段训练代码一行都没执行。原因fastvit的预训练权重托管在外部服务器网络访问超时或证书问题都会导致下载失败。这个错误信息很误导人容易让人以为是模型结构写错了。解决先手动下载权重放到~/.cache/torch/hub/checkpoints/目录下torch加载时会先查本地缓存。或者把create_model里的pretrained临时改成False用随机权重先跑通整个流程确认脚本本身的正确性。等下载好权重后再把pretrained改回True重训。从踩过的坑来说首次跑一个新项目前先检查权重是否已在本地缓存能省掉一晚上的排查时间。4.4 每次训练结果不一样复现性丢失现象相同参数跑两次loss曲线和最终准确率对不上差距还不小。原因随机种子没固定数据加载顺序、随机裁剪、Shuffle、权重初始化都在变。Transformer类的训练对这种随机性敏感不同结果不能稳定对比调参。解决在train.py开头统一固定seeddef set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic True固定seed后同参数跑两次可复现但跨epoch的shuffle顺序仍会变化这是DataLoader的正常行为。cudnn.deterministicTrue会牺牲一点训练速度换取卷积计算的确定性追求复现时值得开。如果还要更严格的可复现把DataLoader的num_workers设成0也可以代价是预处理速度变慢。4.5 导出模型推理结果和训练时不一致现象训练时在验证集上准确率90%导出后用同一张图跑推理结果和训练时打印的不一致有时连类别都变了。原因导出前没调用model.eval()。模型还停留在训练模式BatchNorm在拿当前batch的统计量计算Dropout还在随机丢弃推理输出自然不稳定。解决导出前和推理前都先model.eval()再包torch.no_grad()。如果保存的是加载了state_dict的模型也要确认加载完再切eval顺序不能反load_state_dict不会自动改模式。导出ONNX后建议用onnxruntime验证一下输出差异差异在1e-5以内算正常超过1e-2基本就是eval/训练模式没切换对。5. 导出与测试export_model.py和test.py的完整链路5.1 导出模型格式选择与实现训练结束后export_model.py把训练产物转成可部署的格式。常见做法是两种保存PyTorch的state_dict包或者导出ONNX。state_dict适合继续在PyTorch生态里使用ONNX适合跨框架推理和服务端部署。import torch from fastvit import create_model model create_model(fastvit_t8, pretrainedTrue, num_classes2) model.load_state_dict(torch.load(best_model.pth)) model.eval() # 保存带配置信息的state_dict torch.save({ state_dict: model.state_dict(), num_classes: 2, input_size: (3, 224, 224), arch: fastvit_t8 }, fastvit_export.pt) # 导出onnx trace_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, trace_input, fastvit_t8.onnx, input_names[images], output_names[logits], dynamic_axes{images: {0: batch}, logits: {0: batch}} )保存state_dict时把arch和num_classes一起存进去是我很早就养成的习惯。假如过了半个月要重新部署这个模型只靠best_model.pth完全不知道当时用的什么架构翻聊天记录和笔记都很痛苦。带配置一起保存加载时create_model的参数就有据可依。注意导出ONNX前不要忘记先model.eval()这一步决定了导出后的模型在推理时的行为是否稳定。导出ONNX时dynamic_axes声明了batch维度是动态的这样导出的模型可以接受任意batch size输入。如果不需要动态batch去掉dynamic_axes、固定输入形状反而能获得更好的推理性能。5.2 test.py的推理流程与类别映射test.py做的事情是把导出模型加载回来对图片做预测。核心流程是读图、预处理、前向传播、解析结果。import json import torch from PIL import Image from fastvit import create_model from torchvision import transforms with open(class.json, r) as f: class_map json.load(f) model create_model(fastvit_t8, pretrainedTrue, num_classeslen(class_map)) ckpt torch.load(fastvit_export.pt, map_locationcpu) model.load_state_dict(ckpt[state_dict]) model.eval() transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img Image.open(test.png).convert(RGB) input_tensor transform(img).unsqueeze(0) with torch.no_grad(): logits model(input_tensor) probs torch.softmax(logits, dim1) pred_idx torch.argmax(logits, dim1).item() print(f预测类别: {class_map[str(pred_idx)]}, 置信度: {probs[0][pred_idx].item():.4f})两个细节值得专门讲。第一map_locationcpu让训练时用GPU的权重也能在纯CPU机器上加载部署环境没有CUDA时不会加载失败。第二class_map的key是字符串因为json的key只能是字符串而argmax返回的是int查表前必须str(pred_idx)。我第一次写的时候用整数直接查报KeyError查了半天才意识到类型问题。如果class.json里存的是列表不是字典索引直接用整数就行读之前先看两眼结构再写代码。5.3 result.json和class.json怎么关联result.json记录训练过程指标class.json记录类别映射两者是这条流程里最容易混淆的两个文件。result.json长这样[ {epoch: 1, train_loss: 0.8342, val_acc: 0.6850}, {epoch: 2, train_loss: 0.5123, val_acc: 0.7820} ]class.json长这样{ 0: cat, 1: dog }读取result.json直接用pandasimport pandas as pd df pd.read_json(result.json) print(df.head())预测时需要把result.json里val_acc最高的epoch对应的模型拿出来配合class.json把预测索引翻译成类别名。有一个实用习惯训练脚本里每次在验证集上超越best_acc就保存模型同时把当时的val_acc和epoch记下来这样result.json里的best_acc记录永远和磁盘上的best_model.pth对应不会出现模型文件更新了、日志还停留在上一个epoch的情况。这个对应关系一旦错位调参时会被误导。6. 迁移到自己的数据集三步走完FastVIT迁移把FastVIT这套流程迁移到自己的数据集最直接的应用就是换数据目录、改类别数、重跑脚本。假如要做森林图像分类把原始图片按类别放进子目录makedata.py自动完成训练集和验证集切分train.py只需改num_classesclass.json会在第一次运行后生成新的映射。迁移前建议先跑一遍小数据验证。取每个类别20张图片用t8变体训练5个epoch确认Loss能下降、准确率不是0。这一步能提前暴露数据路径、类别映射、环境依赖的问题等全量数据跑的时候就不会在同一个坑里浪费时间了。迁移学习的参数调整有个常见误区认为换了数据集就要重训全部参数。实际做法是保留ImageNet的预训练权重因为它的浅层特征对大多数自然图像都有效。如果新数据集和ImageNet分布差异很大可以适当调小学习率比如从3e-4改到1e-4让模型在新数据上小步微调避免一下破坏已经学好的特征。我习惯在训练前打印一行模型分类头的结构确认输出维度和class.json的类别数一致print(model.head if hasattr(model, head) else model.classifier)这个习惯帮我躲过两次改完数据集忘记改num_classes的bug那个报错信息相当不直观。现在每次换数据集我都强制走一遍先检查class.json再打印分类头最后跑makedata → train → export → test的完整流程。希望帮到你。本文还有配套的精品资源点击获取