很多刚接触视觉算法的同学都会遇到一个类似困境在文章或课程里把 AlexNet、VGG、ResNet、ViT 的结构图看得明明白白真正打开编辑器准备训练时却发现自己只有“网络结构代码”没有一套完整的数据加载、训练循环、评估和导出脚本。于是从 A 仓库复制数据增强从 B 仓库找训练参数再从 C 仓库拼接模型调用最后训练到一半才发现学习率策略和优化器根本不匹配。这不是个人能力问题而是当前视觉入门教程普遍存在的空白结构讲解很多工程闭环太少。可以给出一个明确判断在图像分类这个方向上从经典 CNN 到 Transformer训练管线其实早已收敛。数据加载、损失函数、优化器、评估逻辑、模型导出这些环节几乎是固定的真正变化的只有模型自身以及少数几处超参数。因此完全可以用一份统一代码模板在 AlexNet、VGG、ResNet、ViT 四个主流架构之间横向切换。这篇文章会围绕一个可运行的分类项目展开用 PyTorch 实现一个统一模型构建层用同一套训练脚本跑通四种架构再把训练好的模型导出为 ONNX 做部署验证。读完你会拥有一份可以复用的视觉分类模板后续换成自己的数据和模型结构只需要改少量代码。1. 这篇文章真正要解决的问题在进入代码之前先把问题说清楚。很多人学习视觉模型时倾向于“每个模型单独学一遍”但这样会陷入一个误区把大量的时间和注意力放在结构代码的差异上而忽略了训练流程的共性。AlexNet 和 ViT 虽然看起来差别巨大但它们都可以塞进一个标准的监督分类流程里读取图像、前向传播、计算交叉熵损失、反向传播、更新参数、周期性评估准确率。这个模板要解决的三个核心问题是如何不写四份几乎一模一样的训练脚本只维护一份训练代码就完成多架构实验。如何在四种架构之间做横向对比而不是每次都在不同的代码风格里来回切换。如何在训练完成后把模型导出为通用格式ONNX让部署代码和训练代码解耦。对读者来说这篇文章的技术收益是掌握一个“模型层可插拔、训练层固定、导出层统一”的最小工程范式。如果你是在校学生可以用它快速完成模型对比实验如果你是算法工程师可以用它作为新项目的主干脚本如果你刚入门深度学习可以通过这份模板搞清楚一条完整训练链路到底由哪些环节组成。这里要先做一个提醒模板的价值在于“统一”但不等于“无脑”。四个架构在超参数敏感度上仍然有明显差异这一点我会在后面的章节专门说明。代码统一解决的是工程重复问题调参经验解决的是模型差异问题两者并不冲突。2. 四个架构并不需要四套训练代码先把四个架构放在一起对比这是理解整篇文章的基础。模型提出年份核心思想典型结构更适用的场景AlexNet2012深度卷积 ReLU Dropout5 层卷积 3 层全连接学习基线、新手入门VGG2014小卷积核反复堆叠多个卷积块 全连接结构直观、迁移学习基线ResNet2015残差连接缓解网络退化基本块 / 瓶颈块通用视觉任务首选ViT2020图像切块 自注意力Patch Embedding Transformer Encoder大数据 预训练微调AlexNet 的历史意义在于它证明了“深度”的价值更深的网络、非线性的 ReLU 激活、数据增强和 Dropout 共同把 ImageNet 分类准确率提升到了一个此前难以想象的高度。虽然今天直接在生产环境中使用 AlexNet 的团队已经不多但它仍然是理解卷积神经网络基本组成的最佳起点。VGG 的核心贡献是简化了卷积核的设计全部使用 3x3 小卷积核通过堆叠卷积块来增加深度。它把“如何设计卷积层”变成了“如何堆叠卷基层”这个思路对后来很多网络都有影响。缺点也明显参数多、计算量大训练和推理成本偏高。ResNet 解决了深度网络难以训练的问题。网络太深时准确率反而会下降这叫“退化问题”。ResNet 通过残差连接让每一层学习的是“残差”——也就是输入和输出之间的差值。这样梯度可以绕过大部分卷积层直接向后传播网络可以做到很深而不容易训练失败。在实际工程中ResNet 系列往往是默认的视觉基线模型。ViT 则完全改变思路把图像切成固定大小的 patch映射成 token然后交给 Transformer 的自注意力机制去建模全局关系。ViT 在超大数据集上训练后再微调效果可以超过 CNN但在数据量不足时它不如 CNN 稳。理解 ViT 的核心不需要写完整代码因为主流框架已经提供官方实现你需要做的是把它接入统一训练流程。从工程视角看这四类架构的差异集中在“特征提取器”和“分类头结构”上而训练上下文完全一致。真正应该掌握的判断是学习重点要分成两层——结构层看差异工程层看共性。3. 环境准备与版本说明本文代码基于 PyTorch 生态示例在 Linux 和 Windows 环境都可以运行。为了保证代码不依赖过新或过旧的库下面列出一个比较稳妥的版本区间。python -c import torch; print(torch.__version__)建议使用Python 3.8 及以上。PyTorch 1.12 及以上推荐 2.x。torchvision 0.13 及以上因为自注意力的 ViT 预训练模型在 torchvision 中从 0.13 开始提供。onnx 和 onnxruntime用于导出和部署验证。安装命令pip install torch torchvision onnx onnxruntime tqdm如果你需要 GPU 加速请根据自己机器的 CUDA 版本选择对应的 PyTorch 安装源这一点以 PyTorch 官网给出的安装命令为准。CPU 环境也能运行本文全部代码只是训练速度会明显变慢建议先用小 batch size 做验证。版本问题是最容易踩坑的环节。torchvision 在较新版本中把pretrainedTrue改成了weights...的形式本文代码为了清晰会直接使用新版写法。如果你使用的是旧版本需要把代码中的weightsIMAGENET1K_V1换成pretrainedTrue。这个细节会在常见问题章节继续展开。4. 数据准备先跑通 CIFAR-10本文示例使用 CIFAR-10 数据集。它是视觉分类最常用的教学数据集50000 张训练图片10000 张测试图片共 10 个类别。数据体积小、下载快特别适合用来验证模板是否跑通。需要注意一个两难问题CIFAR-10 原始图片只有 32x32而 AlexNet、VGG、ResNet、ViT 的标准输入通常是 224x224。直接把小图 resize 到大尺寸会损失细节也会增加显存占用但如果完全不 resize某些模型结构可能无法正常工作尤其是依赖固定 patch 数量的 ViT。模板的处理方式很直接统一 resize 到 224x224。这样四个模型接受完全相同的输入代码路径是统一的符合本文“模板化”的目标。需要理解的是从零训练时把小图放大到 224 并不理想所以当你把模板用到真实数据集时应该根据数据原始尺寸调整Resize的数值。下面是数据加载部分的代码。这里只写一份训练集和验证集都使用同一套逻辑区别在于是否做随机数据增强。# 文件路径data_utils.py import torchvision.transforms as transforms def get_transform(train: bool True): normalize transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) if train: return transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), normalize, ]) return transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), normalize, ])训练阶段加入随机水平翻转和颜色抖动是为了提高模型泛化能力。验证阶段不做随机增强只做 resize 和归一化保证评估结果的稳定性。在train.py中加载数据时使用torchvision.datasets.CIFAR10并设置downloadTrue。第一次运行会自动下载数据之后会直接读取本地缓存。数据规模较小所以 DataLoader 的num_workers设为 2 即可过大会在 Windows 上引发额外问题。5. 统一模型构建层用注册表管理四种架构很多人会问为什么不用简单的 if-else 来切换模型if-else 当然能跑但随着模型种类增多代码会越来越混乱。更好的做法是用一个小的注册表模式每个模型由一个 builder 函数负责构建函数名和模型字符串绑定。以后新增模型只需要加一个函数不需要改调用方。下面这段代码是模板的关键部分。它封装了四个模型的构建逻辑并且支持从 ImageNet 预训练权重初始化。# 文件路径models.py import torch.nn as nn import torchvision _MODEL_FACTORY {} def register_model(name: str): def decorator(fn): _MODEL_FACTORY[name] fn return fn return decorator def build_model(name: str, num_classes: int 10, pretrained: bool False): if name not in _MODEL_FACTORY: raise KeyError(f未知模型{name}可选模型 {list(_MODEL_FACTORY.keys())}) return _MODEL_FACTORY[name](num_classesnum_classes, pretrainedpretrained) register_model(alexnet) def build_alexnet(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.alexnet(weightsweights) model.classifier[6] nn.Linear(model.classifier[6].in_features, num_classes) return model register_model(vgg13) def build_vgg13(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.vgg13(weightsweights) model.classifier[6] nn.Linear(model.classifier[6].in_features, num_classes) return model register_model(resnet18) def build_resnet18(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.resnet18(weightsweights) model.fc nn.Linear(model.fc.in_features, num_classes) return model register_model(vit_b_16) def build_vit_b_16(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.vit_b_16(weightsweights) model.heads.head nn.Linear(model.heads.head.in_features, num_classes) return model这段代码里最重要的一层含义是模型构建是“可插拔”的。调用方永远只依赖build_model(name, num_classes, pretrained)至于 name 对应的是卷积网络还是 Transformer调用方不需要关心。这样训练脚本可以彻底固定下来不会因为换了模型而改动主流程。如果你使用的是较老版本的 torchvisionweights参数不存在。更保守的做法是把weights改成pretrained但那样会失去新版本 API 的优势。这里建议优先升级 torchvision而不是为了兼容旧版本把代码写得更复杂。6. 训练循环一份脚本跑通四个模型训练主脚本是整个模板的核心。它做的事情是所有分类任务共用的加载数据、构建模型、定义损失函数和优化器、迭代若干轮、每轮验证并保存最佳权重。下面给出完整训练代码。这是一个最小可运行版本没有引入分布式训练和复杂的日志框架方便理解主干逻辑。# 文件路径train.py import argparse import os import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision.datasets import CIFAR10 from data_utils import get_transform from models import build_model def train_one_epoch(model, dataloader, criterion, optimizer, device): model.train() total_loss, total_correct, total_num 0.0, 0, 0 for images, labels in dataloader: 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() * images.size(0) total_correct (outputs.argmax(dim1) labels).sum().item() total_num images.size(0) return total_loss / total_num, total_correct / total_num torch.no_grad() def evaluate(model, dataloader, criterion, device): model.eval() total_loss, total_correct, total_num 0.0, 0, 0 for images, labels in dataloader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) total_correct (outputs.argmax(dim1) labels).sum().item() total_num images.size(0) return total_loss / total_num, total_correct / total_num def main(): parser argparse.ArgumentParser() parser.add_argument(--model, typestr, defaultresnet18, choices[alexnet, vgg13, resnet18, vit_b_16]) parser.add_argument(--epochs, typeint, default30) parser.add_argument(--batch-size, typeint, default64) parser.add_argument(--lr, typefloat, default1e-2) parser.add_argument(--optimizer, typestr, defaultsgd, choices[sgd, adamw]) parser.add_argument(--pretrained, actionstore_true) parser.add_argument(--device, typestr, defaultcuda if torch.cuda.is_available() else cpu) parser.add_argument(--data-dir, typestr, default./data) parser.add_argument(--save-dir, typestr, default./checkpoints) args parser.parse_args() os.makedirs(args.save_dir, exist_okTrue) device torch.device(args.device) train_loader DataLoader( CIFAR10(rootargs.data_dir, trainTrue, transformget_transform(True), downloadTrue), batch_sizeargs.batch_size, shuffleTrue, num_workers2, pin_memoryTrue ) val_loader DataLoader( CIFAR10(rootargs.data_dir, trainFalse, transformget_transform(False), downloadTrue), batch_sizeargs.batch_size, shuffleFalse, num_workers2, pin_memoryTrue ) model build_model(args.model, num_classes10, pretrainedargs.pretrained).to(device) criterion nn.CrossEntropyLoss() if args.optimizer sgd: optimizer torch.optim.SGD(model.parameters(), lrargs.lr, momentum0.9, weight_decay5e-4) else: optimizer torch.optim.AdamW(model.parameters(), lrargs.lr, weight_decay5e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxargs.epochs) best_acc 0.0 for epoch in range(args.epochs): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() print(fEpoch {epoch 1}/{args.epochs} | ftrain loss {train_loss:.4f} acc {train_acc * 100:.1f}% | fval acc {val_acc * 100:.1f}% | flr {scheduler.get_last_lr()[0]:.2e}) if val_acc best_acc: best_acc val_acc torch.save({ model: model.state_dict(), epoch: epoch, val_acc: best_acc, }, os.path.join(args.save_dir, f{args.model}_best.pt)) print(f训练完成最佳验证准确率{best_acc * 100:.2f}%) print(f权重已保存至{os.path.join(args.save_dir, f{args.model}_best.pt)}) if __name__ __main__: main()这段代码的逻辑并不复杂但有三个地方需要重点看优化器支持 SGD 和 AdamW。ResNet、VGG 使用带动量的 SGD 通常更稳ViT 更推荐 AdamW配合较小学习率。学习率调度使用余弦退火。它能在训练后期逐步降低学习率提升收敛稳定性。每轮结束后立即在验证集上评估并保存验证准确率最高的权重。这样训练中断时不至于丢失最佳结果。如果你想继续跑第二个模型不需要改动任何代码只需换--model参数。训练脚本本身不感知模型内部结构差异这正是模板化的价值所在。7. 运行与验证从命令行到日志判断在开始训练之前先用最小的配置验证模板本身是否工作正常。这里建议先把 batch size 调小比如 16确定数据加载和模型 forward 没有问题再跑正式训练。四个模型分别通过以下命令启动# AlexNet从零训练 python train.py --model alexnet --epochs 30 --batch-size 64 --lr 1e-2 # VGG13从零训练 python train.py --model vgg13 --epochs 30 --batch-size 64 --lr 1e-2 # ResNet18从零训练 python train.py --model resnet18 --epochs 30 --batch-size 64 --lr 1e-2 # ViT使用 AdamW 并且加载 ImageNet 预训练权重推荐 python train.py --model vit_b_16 --epochs 30 --batch-size 64 --lr 1e-4 --optimizer adamw --pretrained如果你机器显存有限可以继续降低 batch size。Batch size 不影响代码结构只影响训练速度和稳定性。正常运行时的输出格式类似这样Epoch 1/30 | train loss 1.8935 acc 31.2% | val acc 47.1% | lr 1.00e-02 Epoch 2/30 | train loss 1.4102 acc 47.5% | val acc 55.8% | lr 9.66e-03如何判断训练是否健康先看 train loss 是否持续下降。如果 train loss 下降但 val acc 不涨说明过拟合如果 train loss 和 val acc 都不动大概率是学习率设置有问题或者数据加载异常。不同模型在这个模板上的表现会遵循一些规律。ResNet18 从零训练到 30 轮验证准确率通常能接近 80% 甚至更高比较稳定。AlexNet 参数多但结构相对浅在 CIFAR-10 这种小图上提升较慢。VGG13 训练速度慢显存占用高但准确率通常不错。ViT 如果不做预训练在 CIFAR-10 上很容易过拟合这也是为什么它在小数据集上口碑不如 CNN 的直接原因。穿针引线地说模板让你快速得到结果但理解不同模型的行为差异仍然要靠实验对比和调参经验。8. 导出 ONNX训练与部署的解耦模型训练完成后下一步通常是部署。很多人直接把 PyTorch 模型放到服务里这样虽然能用但推理代码和训练代码耦合很重。更通用的做法是导出为 ONNX再用 ONNX Runtime 这种专门的推理引擎加载。ONNX 的好处是与框架无关后续在 CPU、GPU 甚至边缘设备上都能找到对应的运行环境。下面这段代码把训练好的 checkpoint 导出为 ONNX 文件。# 文件路径export_onnx.py import argparse import torch from models import build_model def main(): parser argparse.ArgumentParser() parser.add_argument(--model, typestr, defaultresnet18, choices[alexnet, vgg13, resnet18, vit_b_16]) parser.add_argument(--weight, typestr, requiredTrue) parser.add_argument(--output, typestr, defaultmodel.onnx) parser.add_argument(--size, typeint, default224) args parser.parse_args() model build_model(args.model, num_classes10, pretrainedFalse) state_dict torch.load(args.weight, map_locationcpu) model.load_state_dict(state_dict[model]) model.eval() dummy_input torch.randn(1, 3, args.size, args.size) torch.onnx.export( model, dummy_input, args.output, input_names[input], output_names[logits], dynamic_axes{input: {0: batch}, logits: {0: batch}}, opset_version17, ) print(fONNX 导出完成{args.output}) if __name__ __main__: main()导出命令python export_onnx.py --model resnet18 --weight checkpoints/resnet18_best.pt --output resnet18.onnx导出 ONNX 时最关键的是dynamic_axes。如果不设置ONNX 模型输入 batch 会被固定为 1部署时只能一次推理一张图。设置dynamic_axes后batch 维度可以是动态的更加灵活。导出完成后用 ONNX Runtime 做一次推理验证确认导出没有损坏模型结构。# 文件路径onnx_infer.py import numpy as np import onnxruntime as ort sess ort.InferenceSession(resnet18.onnx) input_name sess.get_inputs()[0].name x np.random.randn(1, 3, 224, 224).astype(np.float32) logits sess.run(None, {input_name: x})[0] print(输出 shape:, logits.shape) print(预测类别索引:, logits.argmax(axis1))这一步的输出应该是一个长度为 10 的向量因为 CIFAR-10 有 10 个类别。如果你导出的模型分类头是 1000 类而不是 10 类说明导出脚本构建模型时没有替换分类头或者加载的权重和模型结构不匹配。这里也能看出模板设计的另一个收益训练权重文件里保存的是“模型参数”而不是完整模型对象导出脚本只依赖build_model这个名字来重建结构。所以模型定义层必须保持稳定这也是前面使用注册表模式的原因。9. 常见问题与排查思路在实际运行中读者最常遇到的并不是模型理论问题而是环境、版本、显存这类工程问题。下面整理了一份排查表按优先级排列。问题现象可能原因排查方式解决方案导入vit_b_16报 ImportErrortorchvision 版本过低没有 ViT 实现执行python -c import torchvision; print(torchvision.__version__)升级 torchvision 到 0.13 以上weights参数不被识别代码使用了新版 API旧版只支持pretrainedTrue查看安装的 PyTorch 版本对应文档改为pretrainedTrue或升级版本显存溢出OOMbatch size 过大或 224x224 输入在 ViT 上显存消耗高观察日志中哪个阶段报错降低 batch size 重试--batch-size 16或更小并减少num_workerstrain loss 不下降学习率过大或过小也可能是数据没有正确归一化打日志检查 loss 数量级尝试 lr 1e-3从零训练用 SGD lr1e-2预训练微调用 lr1e-3ViT 在 CIFAR-10 上效果差小数据量 无预训练Transformer 容易过拟合观察 val acc 是否远低于 train acc使用--pretrained或换用 ResNet 作为基线ONNX 导出后 batch 固定为 1dynamic_axes没有正确设置用 Netron 打开 ONNX 文件查看输入 shape在torch.onnx.export中设置dynamic_axes加载权重时 shape 不匹配checkpoint 的模型和当前--model不一致对比训练命令和导出命令中的模型名确保训练和导出使用同一个模型名称验证准确率高于训练准确率Dropout 和 BatchNorm 在 train/eval 模式行为不同检查model.train()和model.eval()是否正确切换在推理和验证前必须执行model.eval()这张表里的多数问题都无法靠“重装环境”解决而是源于对版本差异和模型行为的不熟悉。建议你在开始训练前先确认三件事torchvision 版本足够新数据集能够正常下载CIFAR-10 的归一化参数已经配置正确。10. 工程实践建议与超参选择经验模板能让你快速把所有模型跑起来但在真实项目中使用时还需要遵守一些工程经验否则很容易被表面上的“跑通”欺骗。第一先小后大。任何新的数据集接入模板都不要直接全量训练。先用每个类别 100 到 200 张图的子集跑 1 到 2 个 epoch确认数据加载正确、loss 能下降、模型可以收敛到随机水平以上。工程里大量 bug 都是在全量训练跑了几小时后才暴露的小规模验证成本低得多。第二学习率策略要区分模型。ResNet 和 VGG 使用 SGD momentum 是经过大量实践验证的稳妥选择学习率从 1e-2 开始调。ViT 更推荐 AdamW配合 warmup 和较小的学习率比如 1e-4 到 5e-4。如果套用同一个学习率到所有模型不是模板的错而是没有理解模型差异的体现。第三数据处理比模型结构更影响最终结果。CIFAR-10 因为图小简单增强就够用但真实数据往往需要做 resize、随机裁剪、旋转、亮度扰动等多种组合。当模板迁移到新数据时首先要调整的是get_transform而不是急着换更复杂的模型。第四checkpoint 里不仅保存模型权重还可以保存优化器状态、当前 epoch、学习率等信息。本文示例为了简洁只保存了模型权重和验证准确率工程中建议至少加上优化器状态便于中断后从最近的状态继续训练。第五从零训练还是微调预训练权重必须想清楚。大数据量任务可以从零训练但视觉任务普遍推荐使用 ImageNet 预训练权重做初始化尤其是 ViT。预训练初始化相当于让模型从“已经会看边的表示”开始而不是从“纯随机”开始这在中小数据集上差别非常大。第六部署前要检查eval()状态。PyTorch 模型在导出 ONNX 前一般要求调用model.eval()因为 BatchNorm 和 Dropout 在训练与推理时的行为不同。如果不做导出的模型行为可能不稳定。11. 总结与后续学习方向这份模板的核心价值不在于代码技巧而在于把视觉分类任务的工程链路固化下来数据、模型、训练、评估、导出。用同一份代码跑 AlexNet、VGG、ResNet、ViT本质上是承认一个事实——这些模型的差异在结构设计层面而工程编排层面已经标准化。理解了这一点你就可以把精力放到真正值得研究的地方。下一步建议从这个方向开始实践先把模板在 CIFAR-10 上完整跑一遍重点观察 ResNet 和 ViT 在相同数据下的行为差异然后用自己的数据集替换 CIFAR-10按实际分类数量修改build_model中的num_classes。如果数据量较大可以尝试使用混合精度训练如果模型较大可以考虑分布式训练如果涉及目标检测或语义分割则需要在本文模板基础上增加对应任务头。最后提醒一句模板是有适用边界的。分类任务的训练管线收敛了不代表所有视觉任务都适合这个套路。当你进入检测、分割、多模态等方向时可以复用本文“注册表管理模型、固定训练循环、统一导出格式”的思维方式但代码结构需要重新设计。把这份模板收藏起来作为你快速验证模型和建立实验基线的起点它会帮你省下大量重复造轮子的时间。
很多刚接触视觉算法的同学都会遇到一个类似困境在文章或课程里把 AlexNet、VGG、ResNet、ViT 的结构图看得明明白白真正打开编辑器准备训练时却发现自己只有“网络结构代码”没有一套完整的数据加载、训练循环、评估和导出脚本。于是从 A 仓库复制数据增强从 B 仓库找训练参数再从 C 仓库拼接模型调用最后训练到一半才发现学习率策略和优化器根本不匹配。这不是个人能力问题而是当前视觉入门教程普遍存在的空白结构讲解很多工程闭环太少。可以给出一个明确判断在图像分类这个方向上从经典 CNN 到 Transformer训练管线其实早已收敛。数据加载、损失函数、优化器、评估逻辑、模型导出这些环节几乎是固定的真正变化的只有模型自身以及少数几处超参数。因此完全可以用一份统一代码模板在 AlexNet、VGG、ResNet、ViT 四个主流架构之间横向切换。这篇文章会围绕一个可运行的分类项目展开用 PyTorch 实现一个统一模型构建层用同一套训练脚本跑通四种架构再把训练好的模型导出为 ONNX 做部署验证。读完你会拥有一份可以复用的视觉分类模板后续换成自己的数据和模型结构只需要改少量代码。1. 这篇文章真正要解决的问题在进入代码之前先把问题说清楚。很多人学习视觉模型时倾向于“每个模型单独学一遍”但这样会陷入一个误区把大量的时间和注意力放在结构代码的差异上而忽略了训练流程的共性。AlexNet 和 ViT 虽然看起来差别巨大但它们都可以塞进一个标准的监督分类流程里读取图像、前向传播、计算交叉熵损失、反向传播、更新参数、周期性评估准确率。这个模板要解决的三个核心问题是如何不写四份几乎一模一样的训练脚本只维护一份训练代码就完成多架构实验。如何在四种架构之间做横向对比而不是每次都在不同的代码风格里来回切换。如何在训练完成后把模型导出为通用格式ONNX让部署代码和训练代码解耦。对读者来说这篇文章的技术收益是掌握一个“模型层可插拔、训练层固定、导出层统一”的最小工程范式。如果你是在校学生可以用它快速完成模型对比实验如果你是算法工程师可以用它作为新项目的主干脚本如果你刚入门深度学习可以通过这份模板搞清楚一条完整训练链路到底由哪些环节组成。这里要先做一个提醒模板的价值在于“统一”但不等于“无脑”。四个架构在超参数敏感度上仍然有明显差异这一点我会在后面的章节专门说明。代码统一解决的是工程重复问题调参经验解决的是模型差异问题两者并不冲突。2. 四个架构并不需要四套训练代码先把四个架构放在一起对比这是理解整篇文章的基础。模型提出年份核心思想典型结构更适用的场景AlexNet2012深度卷积 ReLU Dropout5 层卷积 3 层全连接学习基线、新手入门VGG2014小卷积核反复堆叠多个卷积块 全连接结构直观、迁移学习基线ResNet2015残差连接缓解网络退化基本块 / 瓶颈块通用视觉任务首选ViT2020图像切块 自注意力Patch Embedding Transformer Encoder大数据 预训练微调AlexNet 的历史意义在于它证明了“深度”的价值更深的网络、非线性的 ReLU 激活、数据增强和 Dropout 共同把 ImageNet 分类准确率提升到了一个此前难以想象的高度。虽然今天直接在生产环境中使用 AlexNet 的团队已经不多但它仍然是理解卷积神经网络基本组成的最佳起点。VGG 的核心贡献是简化了卷积核的设计全部使用 3x3 小卷积核通过堆叠卷积块来增加深度。它把“如何设计卷积层”变成了“如何堆叠卷基层”这个思路对后来很多网络都有影响。缺点也明显参数多、计算量大训练和推理成本偏高。ResNet 解决了深度网络难以训练的问题。网络太深时准确率反而会下降这叫“退化问题”。ResNet 通过残差连接让每一层学习的是“残差”——也就是输入和输出之间的差值。这样梯度可以绕过大部分卷积层直接向后传播网络可以做到很深而不容易训练失败。在实际工程中ResNet 系列往往是默认的视觉基线模型。ViT 则完全改变思路把图像切成固定大小的 patch映射成 token然后交给 Transformer 的自注意力机制去建模全局关系。ViT 在超大数据集上训练后再微调效果可以超过 CNN但在数据量不足时它不如 CNN 稳。理解 ViT 的核心不需要写完整代码因为主流框架已经提供官方实现你需要做的是把它接入统一训练流程。从工程视角看这四类架构的差异集中在“特征提取器”和“分类头结构”上而训练上下文完全一致。真正应该掌握的判断是学习重点要分成两层——结构层看差异工程层看共性。3. 环境准备与版本说明本文代码基于 PyTorch 生态示例在 Linux 和 Windows 环境都可以运行。为了保证代码不依赖过新或过旧的库下面列出一个比较稳妥的版本区间。python -c import torch; print(torch.__version__)建议使用Python 3.8 及以上。PyTorch 1.12 及以上推荐 2.x。torchvision 0.13 及以上因为自注意力的 ViT 预训练模型在 torchvision 中从 0.13 开始提供。onnx 和 onnxruntime用于导出和部署验证。安装命令pip install torch torchvision onnx onnxruntime tqdm如果你需要 GPU 加速请根据自己机器的 CUDA 版本选择对应的 PyTorch 安装源这一点以 PyTorch 官网给出的安装命令为准。CPU 环境也能运行本文全部代码只是训练速度会明显变慢建议先用小 batch size 做验证。版本问题是最容易踩坑的环节。torchvision 在较新版本中把pretrainedTrue改成了weights...的形式本文代码为了清晰会直接使用新版写法。如果你使用的是旧版本需要把代码中的weightsIMAGENET1K_V1换成pretrainedTrue。这个细节会在常见问题章节继续展开。4. 数据准备先跑通 CIFAR-10本文示例使用 CIFAR-10 数据集。它是视觉分类最常用的教学数据集50000 张训练图片10000 张测试图片共 10 个类别。数据体积小、下载快特别适合用来验证模板是否跑通。需要注意一个两难问题CIFAR-10 原始图片只有 32x32而 AlexNet、VGG、ResNet、ViT 的标准输入通常是 224x224。直接把小图 resize 到大尺寸会损失细节也会增加显存占用但如果完全不 resize某些模型结构可能无法正常工作尤其是依赖固定 patch 数量的 ViT。模板的处理方式很直接统一 resize 到 224x224。这样四个模型接受完全相同的输入代码路径是统一的符合本文“模板化”的目标。需要理解的是从零训练时把小图放大到 224 并不理想所以当你把模板用到真实数据集时应该根据数据原始尺寸调整Resize的数值。下面是数据加载部分的代码。这里只写一份训练集和验证集都使用同一套逻辑区别在于是否做随机数据增强。# 文件路径data_utils.py import torchvision.transforms as transforms def get_transform(train: bool True): normalize transforms.Normalize( mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) if train: return transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), normalize, ]) return transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), normalize, ])训练阶段加入随机水平翻转和颜色抖动是为了提高模型泛化能力。验证阶段不做随机增强只做 resize 和归一化保证评估结果的稳定性。在train.py中加载数据时使用torchvision.datasets.CIFAR10并设置downloadTrue。第一次运行会自动下载数据之后会直接读取本地缓存。数据规模较小所以 DataLoader 的num_workers设为 2 即可过大会在 Windows 上引发额外问题。5. 统一模型构建层用注册表管理四种架构很多人会问为什么不用简单的 if-else 来切换模型if-else 当然能跑但随着模型种类增多代码会越来越混乱。更好的做法是用一个小的注册表模式每个模型由一个 builder 函数负责构建函数名和模型字符串绑定。以后新增模型只需要加一个函数不需要改调用方。下面这段代码是模板的关键部分。它封装了四个模型的构建逻辑并且支持从 ImageNet 预训练权重初始化。# 文件路径models.py import torch.nn as nn import torchvision _MODEL_FACTORY {} def register_model(name: str): def decorator(fn): _MODEL_FACTORY[name] fn return fn return decorator def build_model(name: str, num_classes: int 10, pretrained: bool False): if name not in _MODEL_FACTORY: raise KeyError(f未知模型{name}可选模型 {list(_MODEL_FACTORY.keys())}) return _MODEL_FACTORY[name](num_classesnum_classes, pretrainedpretrained) register_model(alexnet) def build_alexnet(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.alexnet(weightsweights) model.classifier[6] nn.Linear(model.classifier[6].in_features, num_classes) return model register_model(vgg13) def build_vgg13(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.vgg13(weightsweights) model.classifier[6] nn.Linear(model.classifier[6].in_features, num_classes) return model register_model(resnet18) def build_resnet18(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.resnet18(weightsweights) model.fc nn.Linear(model.fc.in_features, num_classes) return model register_model(vit_b_16) def build_vit_b_16(num_classes: int, pretrained: bool False): weights IMAGENET1K_V1 if pretrained else None model torchvision.models.vit_b_16(weightsweights) model.heads.head nn.Linear(model.heads.head.in_features, num_classes) return model这段代码里最重要的一层含义是模型构建是“可插拔”的。调用方永远只依赖build_model(name, num_classes, pretrained)至于 name 对应的是卷积网络还是 Transformer调用方不需要关心。这样训练脚本可以彻底固定下来不会因为换了模型而改动主流程。如果你使用的是较老版本的 torchvisionweights参数不存在。更保守的做法是把weights改成pretrained但那样会失去新版本 API 的优势。这里建议优先升级 torchvision而不是为了兼容旧版本把代码写得更复杂。6. 训练循环一份脚本跑通四个模型训练主脚本是整个模板的核心。它做的事情是所有分类任务共用的加载数据、构建模型、定义损失函数和优化器、迭代若干轮、每轮验证并保存最佳权重。下面给出完整训练代码。这是一个最小可运行版本没有引入分布式训练和复杂的日志框架方便理解主干逻辑。# 文件路径train.py import argparse import os import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision.datasets import CIFAR10 from data_utils import get_transform from models import build_model def train_one_epoch(model, dataloader, criterion, optimizer, device): model.train() total_loss, total_correct, total_num 0.0, 0, 0 for images, labels in dataloader: 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() * images.size(0) total_correct (outputs.argmax(dim1) labels).sum().item() total_num images.size(0) return total_loss / total_num, total_correct / total_num torch.no_grad() def evaluate(model, dataloader, criterion, device): model.eval() total_loss, total_correct, total_num 0.0, 0, 0 for images, labels in dataloader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) total_loss loss.item() * images.size(0) total_correct (outputs.argmax(dim1) labels).sum().item() total_num images.size(0) return total_loss / total_num, total_correct / total_num def main(): parser argparse.ArgumentParser() parser.add_argument(--model, typestr, defaultresnet18, choices[alexnet, vgg13, resnet18, vit_b_16]) parser.add_argument(--epochs, typeint, default30) parser.add_argument(--batch-size, typeint, default64) parser.add_argument(--lr, typefloat, default1e-2) parser.add_argument(--optimizer, typestr, defaultsgd, choices[sgd, adamw]) parser.add_argument(--pretrained, actionstore_true) parser.add_argument(--device, typestr, defaultcuda if torch.cuda.is_available() else cpu) parser.add_argument(--data-dir, typestr, default./data) parser.add_argument(--save-dir, typestr, default./checkpoints) args parser.parse_args() os.makedirs(args.save_dir, exist_okTrue) device torch.device(args.device) train_loader DataLoader( CIFAR10(rootargs.data_dir, trainTrue, transformget_transform(True), downloadTrue), batch_sizeargs.batch_size, shuffleTrue, num_workers2, pin_memoryTrue ) val_loader DataLoader( CIFAR10(rootargs.data_dir, trainFalse, transformget_transform(False), downloadTrue), batch_sizeargs.batch_size, shuffleFalse, num_workers2, pin_memoryTrue ) model build_model(args.model, num_classes10, pretrainedargs.pretrained).to(device) criterion nn.CrossEntropyLoss() if args.optimizer sgd: optimizer torch.optim.SGD(model.parameters(), lrargs.lr, momentum0.9, weight_decay5e-4) else: optimizer torch.optim.AdamW(model.parameters(), lrargs.lr, weight_decay5e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxargs.epochs) best_acc 0.0 for epoch in range(args.epochs): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer, device) val_loss, val_acc evaluate(model, val_loader, criterion, device) scheduler.step() print(fEpoch {epoch 1}/{args.epochs} | ftrain loss {train_loss:.4f} acc {train_acc * 100:.1f}% | fval acc {val_acc * 100:.1f}% | flr {scheduler.get_last_lr()[0]:.2e}) if val_acc best_acc: best_acc val_acc torch.save({ model: model.state_dict(), epoch: epoch, val_acc: best_acc, }, os.path.join(args.save_dir, f{args.model}_best.pt)) print(f训练完成最佳验证准确率{best_acc * 100:.2f}%) print(f权重已保存至{os.path.join(args.save_dir, f{args.model}_best.pt)}) if __name__ __main__: main()这段代码的逻辑并不复杂但有三个地方需要重点看优化器支持 SGD 和 AdamW。ResNet、VGG 使用带动量的 SGD 通常更稳ViT 更推荐 AdamW配合较小学习率。学习率调度使用余弦退火。它能在训练后期逐步降低学习率提升收敛稳定性。每轮结束后立即在验证集上评估并保存验证准确率最高的权重。这样训练中断时不至于丢失最佳结果。如果你想继续跑第二个模型不需要改动任何代码只需换--model参数。训练脚本本身不感知模型内部结构差异这正是模板化的价值所在。7. 运行与验证从命令行到日志判断在开始训练之前先用最小的配置验证模板本身是否工作正常。这里建议先把 batch size 调小比如 16确定数据加载和模型 forward 没有问题再跑正式训练。四个模型分别通过以下命令启动# AlexNet从零训练 python train.py --model alexnet --epochs 30 --batch-size 64 --lr 1e-2 # VGG13从零训练 python train.py --model vgg13 --epochs 30 --batch-size 64 --lr 1e-2 # ResNet18从零训练 python train.py --model resnet18 --epochs 30 --batch-size 64 --lr 1e-2 # ViT使用 AdamW 并且加载 ImageNet 预训练权重推荐 python train.py --model vit_b_16 --epochs 30 --batch-size 64 --lr 1e-4 --optimizer adamw --pretrained如果你机器显存有限可以继续降低 batch size。Batch size 不影响代码结构只影响训练速度和稳定性。正常运行时的输出格式类似这样Epoch 1/30 | train loss 1.8935 acc 31.2% | val acc 47.1% | lr 1.00e-02 Epoch 2/30 | train loss 1.4102 acc 47.5% | val acc 55.8% | lr 9.66e-03如何判断训练是否健康先看 train loss 是否持续下降。如果 train loss 下降但 val acc 不涨说明过拟合如果 train loss 和 val acc 都不动大概率是学习率设置有问题或者数据加载异常。不同模型在这个模板上的表现会遵循一些规律。ResNet18 从零训练到 30 轮验证准确率通常能接近 80% 甚至更高比较稳定。AlexNet 参数多但结构相对浅在 CIFAR-10 这种小图上提升较慢。VGG13 训练速度慢显存占用高但准确率通常不错。ViT 如果不做预训练在 CIFAR-10 上很容易过拟合这也是为什么它在小数据集上口碑不如 CNN 的直接原因。穿针引线地说模板让你快速得到结果但理解不同模型的行为差异仍然要靠实验对比和调参经验。8. 导出 ONNX训练与部署的解耦模型训练完成后下一步通常是部署。很多人直接把 PyTorch 模型放到服务里这样虽然能用但推理代码和训练代码耦合很重。更通用的做法是导出为 ONNX再用 ONNX Runtime 这种专门的推理引擎加载。ONNX 的好处是与框架无关后续在 CPU、GPU 甚至边缘设备上都能找到对应的运行环境。下面这段代码把训练好的 checkpoint 导出为 ONNX 文件。# 文件路径export_onnx.py import argparse import torch from models import build_model def main(): parser argparse.ArgumentParser() parser.add_argument(--model, typestr, defaultresnet18, choices[alexnet, vgg13, resnet18, vit_b_16]) parser.add_argument(--weight, typestr, requiredTrue) parser.add_argument(--output, typestr, defaultmodel.onnx) parser.add_argument(--size, typeint, default224) args parser.parse_args() model build_model(args.model, num_classes10, pretrainedFalse) state_dict torch.load(args.weight, map_locationcpu) model.load_state_dict(state_dict[model]) model.eval() dummy_input torch.randn(1, 3, args.size, args.size) torch.onnx.export( model, dummy_input, args.output, input_names[input], output_names[logits], dynamic_axes{input: {0: batch}, logits: {0: batch}}, opset_version17, ) print(fONNX 导出完成{args.output}) if __name__ __main__: main()导出命令python export_onnx.py --model resnet18 --weight checkpoints/resnet18_best.pt --output resnet18.onnx导出 ONNX 时最关键的是dynamic_axes。如果不设置ONNX 模型输入 batch 会被固定为 1部署时只能一次推理一张图。设置dynamic_axes后batch 维度可以是动态的更加灵活。导出完成后用 ONNX Runtime 做一次推理验证确认导出没有损坏模型结构。# 文件路径onnx_infer.py import numpy as np import onnxruntime as ort sess ort.InferenceSession(resnet18.onnx) input_name sess.get_inputs()[0].name x np.random.randn(1, 3, 224, 224).astype(np.float32) logits sess.run(None, {input_name: x})[0] print(输出 shape:, logits.shape) print(预测类别索引:, logits.argmax(axis1))这一步的输出应该是一个长度为 10 的向量因为 CIFAR-10 有 10 个类别。如果你导出的模型分类头是 1000 类而不是 10 类说明导出脚本构建模型时没有替换分类头或者加载的权重和模型结构不匹配。这里也能看出模板设计的另一个收益训练权重文件里保存的是“模型参数”而不是完整模型对象导出脚本只依赖build_model这个名字来重建结构。所以模型定义层必须保持稳定这也是前面使用注册表模式的原因。9. 常见问题与排查思路在实际运行中读者最常遇到的并不是模型理论问题而是环境、版本、显存这类工程问题。下面整理了一份排查表按优先级排列。问题现象可能原因排查方式解决方案导入vit_b_16报 ImportErrortorchvision 版本过低没有 ViT 实现执行python -c import torchvision; print(torchvision.__version__)升级 torchvision 到 0.13 以上weights参数不被识别代码使用了新版 API旧版只支持pretrainedTrue查看安装的 PyTorch 版本对应文档改为pretrainedTrue或升级版本显存溢出OOMbatch size 过大或 224x224 输入在 ViT 上显存消耗高观察日志中哪个阶段报错降低 batch size 重试--batch-size 16或更小并减少num_workerstrain loss 不下降学习率过大或过小也可能是数据没有正确归一化打日志检查 loss 数量级尝试 lr 1e-3从零训练用 SGD lr1e-2预训练微调用 lr1e-3ViT 在 CIFAR-10 上效果差小数据量 无预训练Transformer 容易过拟合观察 val acc 是否远低于 train acc使用--pretrained或换用 ResNet 作为基线ONNX 导出后 batch 固定为 1dynamic_axes没有正确设置用 Netron 打开 ONNX 文件查看输入 shape在torch.onnx.export中设置dynamic_axes加载权重时 shape 不匹配checkpoint 的模型和当前--model不一致对比训练命令和导出命令中的模型名确保训练和导出使用同一个模型名称验证准确率高于训练准确率Dropout 和 BatchNorm 在 train/eval 模式行为不同检查model.train()和model.eval()是否正确切换在推理和验证前必须执行model.eval()这张表里的多数问题都无法靠“重装环境”解决而是源于对版本差异和模型行为的不熟悉。建议你在开始训练前先确认三件事torchvision 版本足够新数据集能够正常下载CIFAR-10 的归一化参数已经配置正确。10. 工程实践建议与超参选择经验模板能让你快速把所有模型跑起来但在真实项目中使用时还需要遵守一些工程经验否则很容易被表面上的“跑通”欺骗。第一先小后大。任何新的数据集接入模板都不要直接全量训练。先用每个类别 100 到 200 张图的子集跑 1 到 2 个 epoch确认数据加载正确、loss 能下降、模型可以收敛到随机水平以上。工程里大量 bug 都是在全量训练跑了几小时后才暴露的小规模验证成本低得多。第二学习率策略要区分模型。ResNet 和 VGG 使用 SGD momentum 是经过大量实践验证的稳妥选择学习率从 1e-2 开始调。ViT 更推荐 AdamW配合 warmup 和较小的学习率比如 1e-4 到 5e-4。如果套用同一个学习率到所有模型不是模板的错而是没有理解模型差异的体现。第三数据处理比模型结构更影响最终结果。CIFAR-10 因为图小简单增强就够用但真实数据往往需要做 resize、随机裁剪、旋转、亮度扰动等多种组合。当模板迁移到新数据时首先要调整的是get_transform而不是急着换更复杂的模型。第四checkpoint 里不仅保存模型权重还可以保存优化器状态、当前 epoch、学习率等信息。本文示例为了简洁只保存了模型权重和验证准确率工程中建议至少加上优化器状态便于中断后从最近的状态继续训练。第五从零训练还是微调预训练权重必须想清楚。大数据量任务可以从零训练但视觉任务普遍推荐使用 ImageNet 预训练权重做初始化尤其是 ViT。预训练初始化相当于让模型从“已经会看边的表示”开始而不是从“纯随机”开始这在中小数据集上差别非常大。第六部署前要检查eval()状态。PyTorch 模型在导出 ONNX 前一般要求调用model.eval()因为 BatchNorm 和 Dropout 在训练与推理时的行为不同。如果不做导出的模型行为可能不稳定。11. 总结与后续学习方向这份模板的核心价值不在于代码技巧而在于把视觉分类任务的工程链路固化下来数据、模型、训练、评估、导出。用同一份代码跑 AlexNet、VGG、ResNet、ViT本质上是承认一个事实——这些模型的差异在结构设计层面而工程编排层面已经标准化。理解了这一点你就可以把精力放到真正值得研究的地方。下一步建议从这个方向开始实践先把模板在 CIFAR-10 上完整跑一遍重点观察 ResNet 和 ViT 在相同数据下的行为差异然后用自己的数据集替换 CIFAR-10按实际分类数量修改build_model中的num_classes。如果数据量较大可以尝试使用混合精度训练如果模型较大可以考虑分布式训练如果涉及目标检测或语义分割则需要在本文模板基础上增加对应任务头。最后提醒一句模板是有适用边界的。分类任务的训练管线收敛了不代表所有视觉任务都适合这个套路。当你进入检测、分割、多模态等方向时可以复用本文“注册表管理模型、固定训练循环、统一导出格式”的思维方式但代码结构需要重新设计。把这份模板收藏起来作为你快速验证模型和建立实验基线的起点它会帮你省下大量重复造轮子的时间。
