FastVIT图像分类实战：线性注意力加速Transformer落地的完整流程

简介：这是一套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', encoding='utf-8') as f: json.dump(class_map, f, ensure_ascii=False, indent=2)

这里我习惯对类别名做一次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_ratio=0.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_ok=True) img.save(os.path.join(out_dir, f))

参数说明几点。val_ratio=0.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(p=0.5), transforms.ColorJitter(brightness=0.2, contrast=0.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的scale=0.8到1.0表示裁剪区域占原图面积的比例。数据集小时可以放宽到0.5，增强强度加大，但太小会裁掉太多上下文，模型反而容易学偏。ColorJitter的亮度、对比度各0.2是温和档，深色或曝光不匀的数据集可以再降一点。Normalize必须用ImageNet的均值方差，因为后面的FastVIT加载的是ImageNet预训练权重，输入分布要和预训练时对齐。如果从头训练，最好按自己的数据集重新统计mean和std，否则开头几个epoch的loss会飘得比较厉害。

数据准备好之后，用torchvision的ImageFolder加载，再用DataLoader组织成batch：

from torch.utils.data import DataLoader from torchvision.datasets import ImageFolder train_dataset = ImageFolder('data/train', transform=train_transform) val_dataset = ImageFolder('data/val', transform=val_transform) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4, pin_memory=True)

num_workers=4是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块的宽度和深度。换变体只需要改一个字符串，不需要动训练脚本里的其他逻辑。

选型时不要盲目上最大的模型。我在小数据集上用m16做过对比，因为过拟合严重，最终准确率反而不如t12。模型容量要跟数据量匹配：每个类别样本少于500张时，优先考虑t8或t12；数据量上万再考虑s12以上。训练中初始化模型的写法：

from fastvit import create_model model = create_model( 'fastvit_t8', pretrained=True, num_classes=2, # 必须和class.json的类别数一致 dropout=0.1 )

pretrained=True会加载ImageNet预训练权重。这个参数对收敛速度影响非常大，小数据集从随机初始化开始训，经常几百个epoch都难收敛，加载预训练后几十个epoch就能到可用精度。dropout=0.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(), lr=3e-4, weight_decay=0.05) scheduler = CosineAnnealingLR(optimizer, T_max=30)

学习率从3e-4起步比较稳，数据量小降到1e-4，数据量大可以尝试5e-4。weight_decay=0.05是ViT训练范式里的经验值，作用对象是权重本身，bias和LayerNorm参数一般不参与。T_max=30表示余弦退火半周期长度，通常和总epoch数一致；如果总epoch是60，T_max设30的话学习率会在30个epoch时降到最低点然后回升，这种周期策略在很多任务上效果更好。

类别不平衡时给交叉熵加权重：

class_weights = torch.tensor([1.0, 5.0]).to(device) criterion = nn.CrossEntropyLoss(weight=class_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, dim=1) correct += (preds == labels).sum().item() total += labels.size(0) acc = correct / total scheduler.step() print(f'Epoch {epoch+1}/{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', encoding='utf-8') as f: json.dump(results, f, indent=2, ensure_ascii=False)

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开头统一固定seed：

def 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.deterministic=True会牺牲一点训练速度换取卷积计算的确定性，追求复现时值得开。如果还要更严格的可复现，把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', pretrained=True, num_classes=2) 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', pretrained=True, num_classes=len(class_map)) ckpt = torch.load('fastvit_export.pt', map_location='cpu') 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, dim=1) pred_idx = torch.argmax(logits, dim=1).item() print(f'预测类别: {class_map[str(pred_idx)]}, 置信度: {probs[0][pred_idx].item():.4f}')

两个细节值得专门讲。第一，map_location='cpu'让训练时用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直接用pandas：

import 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_classes，class.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的完整流程。希望帮到你。

本文还有配套的精品资源，点击获取