4、PyTorch模型训练与导出：使用timm库加载预训练模型、自定义数据集训练、模型评估与保存

好，咱们进入正题。这一章要解决的核心问题很明确：怎么把MobileViT从PyTorch里训练好，再导出成Android能用的格式。我见过太多人卡在这一步——模型在服务器上跑得飞起，一到手机端就各种报错。说白了，就是训练和导出环节没对齐。

我个人习惯把整个流程拆成四块：加载预训练模型 → 准备自定义数据 → 训练调优 → 评估导出。咱们一个一个来。

4.1 用timm库加载MobileViT预训练模型

timm这个库，我愿称之为「模型动物园」。它把主流视觉模型都封装好了，一行代码就能加载。你想想看，要是自己从头写MobileViT，光注意力机制就得调好几天。

import timm
import torch

# 加载MobileViT-Small，预训练权重来自ImageNet-1K
model = timm.create_model('mobilevit_small', pretrained=True)
print(model.default_cfg)  # 看看模型配置

这里有个坑，我曾经踩过：timm里的MobileViT版本和不完全一致。比如输出通道数、下采样倍数，可能差那么一点点。我的建议是，加载后先打印一下模型结构，确认一下。

小技巧： 如果你要部署到低端Android设备，可以用 mobilevit_xxs 版本，参数量只有1.3M，速度能快不少。

4.2 自定义数据集训练——别让数据成为瓶颈

实际项目中，很少直接用ImageNet。你得用自己的数据。我接过一个项目，客户给了10万张工业缺陷图片，但标签格式乱七八糟。嗯，这里要注意，数据加载的效率直接影响训练速度。

我一般这么组织数据集：

dataset/
├── train/
│   ├── class_0/
│   │   ├── img_001.jpg
│   │   └── img_002.jpg
│   └── class_1/
│       └── ...
└── val/
    ├── class_0/
    └── class_1/

然后用 torchvision.datasets.ImageFolder 直接加载，配合 DataLoader 做多线程读取。

from torchvision import transforms, datasets
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.RandomCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

train_dataset = datasets.ImageFolder('dataset/train', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4)
注意： MobileViT的输入尺寸是256x256，但timm默认用224x224。我建议统一用256，因为下采样后特征图尺寸更规整，对移动端部署更友好。

4.3 训练配置——我的一些经验参数

训练MobileViT，说白了就是调三个东西：学习率、优化器、损失函数。我直接给出一套我验证过的配置：

参数 推荐值 说明
优化器 AdamW 比SGD收敛快，适合迁移学习
初始学习率 1e-4 预训练模型用这个，别太大
权重衰减 1e-5 防止过拟合，尤其小数据集
批次大小 64 根据显存调整，16G显存够用
学习率调度 CosineAnnealing 周期下降，效果稳定 
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR

optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-5)
scheduler = CosineAnnealingLR(optimizer, T_max=50)
criterion = torch.nn.CrossEntropyLoss()

为什么会用AdamW而不是Adam？我在项目中对比过，AdamW的权重衰减实现更干净，不会把正则项和梯度更新混在一起。尤其当你用预训练模型时，这个差异会被放大。

4.4 训练循环——别只盯着loss看

很多人训练时只看loss下降就开心了。我建议你同时关注验证集的准确率和损失。loss下降但acc不动？那可能是过拟合了。

def train_one_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()
    
    return total_loss / len(loader), 100. * correct / total

嗯，这里有个细节：MobileViT的推理速度在训练时可能比ResNet慢，因为它的Transformer结构计算量更大。我建议用混合精度训练（AMP）来加速。

核心经验： 对于自定义数据集，我一般先冻结backbone训练分类头50个epoch，再全模型微调30个epoch。这样既快又稳。

4.5 模型评估——用混淆矩阵说话

光看准确率不够。我遇到过准确率95%但某个类别完全没识别出来的情况。所以，混淆矩阵和分类报告是必须的。

from sklearn.metrics import classification_report, confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt

def evaluate(model, loader, device, class_names):
    model.eval()
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for images, labels in loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, preds = outputs.max(1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
    
    print(classification_report(all_labels, all_preds, target_names=class_names))
    
    # 绘制混淆矩阵
    cm = confusion_matrix(all_labels, all_preds)
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=class_names, yticklabels=class_names)
    plt.savefig('confusion_matrix.png')

4.6 模型保存——为导出做准备

保存模型不是简单调个 torch.save 就完事了。你要考虑后续的ONNX导出和量化。我建议保存两种格式：

  1. 完整checkpoint：包含优化器状态，用于继续训练
  2. 推理状态字典：只保留模型权重，用于导出
# 保存完整checkpoint（用于恢复训练）
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'best_acc': best_acc,
}
torch.save(checkpoint, 'mobilevit_checkpoint.pth')

# 保存推理模型（用于导出）
torch.save(model.state_dict(), 'mobilevit_inference.pth')
重要提醒： 保存推理模型前，一定要调用 model.eval()。否则导出的模型里会包含Dropout和BatchNorm的训练状态，部署到Android上推理结果会错。

4.7 本章知识体系

我把整个流程画了张图，方便你理解各环节的关系：

MobileViT训练与导出流程 1. 加载预训练模型 timm.create_model 2. 自定义数据集 ImageFolder + DataLoader 3. 训练调优 AdamW + CosineAnnealing 4. 模型评估 混淆矩阵 + 分类报告 5. 模型保存 checkpoint + inference 迭代调优 最终输出：可用于ONNX/TFLite导出的模型权重

这张图把五个阶段串起来了。你注意看左侧那条虚线——评估结果不满意就回去调参，这是常态。我做过一个项目，光调学习率就试了20组，最后发现是数据增强太强导致模型学不到特征。

我的习惯： 每次训练都记录实验日志，包括学习率、batch size、数据增强参数、最终acc。这样下次遇到类似任务，直接翻日志就能找到最佳配置。

好了，这一章的内容就这些。你按照这个流程走一遍，MobileViT的训练和导出应该不会出大问题。记住，训练是手段，导出到Android才是目的。下一章我们会讲怎么把训练好的模型转成ONNX，再进一步量化成TFLite——那才是真正上手机的关键一步。

可可精品资料库

← 上一章 📖 返回目录 下一章 →