简介这是一份基于PyTorch的mobilenetv2图像分类实战资源包完整覆盖从模型训练到TensorRT部署的端到端流程适合深度学习图像分类与模型加速方向的开发者和学生。包内数据以植物幼苗数据集为基础共包含12个类别文件总计2000个其中2454张PNG图片构成主要训练样本8个Python脚本和2个pyc编译文件覆盖数据集加载、模型搭建、训练验证、预测及onnx/TensorRT转换推理等关键环节压缩包大小约936MB。目前已有1854人学习下载。借助此资源包读者可系统掌握torchvision.models中mobilenetv2的调用方法、自定义数据集的方式以及Cutout和Mixup两种数据增强技巧还能学会使用余弦退火动态调整学习率完成训练和验证并载入训练好的模型进行预测。此资源包还演示了PyTorch模型转ONNX及ONNX转TensorRT的完整流程并分别实现推理帮助读者建立从分类建模到部署加速的整体认知。1. mobilenetv2 图像分类从训练到部署卡在 TensorRT 这一步的看过来mobilenetv2PyTorch 版从训练走到 TensorRT 部署中间隔着的不是代码量而是一堆只有踩过才知道的坑——ONNX 导出的动态轴、TensorRT 对 BN 层的融合、FP16 在某些层上的精度抖动随便碰上一个推理结果就从 95% 准确率变成噪声。这套实战资源把一条完整链路打通了从 PyTorch 环境配置、MobileNetV2 结构复现、迁移学习训练到导出 ONNX、构建 TensorRT 引擎最后给出端到端的性能验证脚本。适合两类人刚学完 PyTorch 基础、想用真实图像分类项目练手的初学者以及模型已经训练完、卡在部署提速环节的工程师。不绕理论按步骤跑通为止。2. PyTorch 环境与 MobileNetV2 模型先让训练链路跑通再谈加速2.1 环境配置CUDA、PyTorch、TensorRT 的版本搭配很多人上来就装最新版结果 PyTorch 能跑、TensorRT 加载引擎报INVALID_ENGINE最后发现是 CUDA 运行时版本和 TensorRT 编译版本对不上。我一般先把版本矩阵定死再动手装。conda create -n mobilenetv2 python3.10 -y conda activate mobilenetv2 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121这里cu121对应 CUDA 12.1 的预编译版本PyTorch 2.x 全系列都支持。装完之后验证一下 GPU 是否真的可用python -c import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0))输出True NVIDIA GeForce RTX 3060这类结果才算通过。如果torch.cuda.is_available()返回False九成是 PyTorch 装了 CPU 版或者 CUDA 驱动版本太老先把 NVIDIA 驱动升到 525 以上再看。TensorRT 这边我建议直接用 NVIDIA 官网下载的 tar 包不要走 apt因为 apt 源里的版本往往和 PyTorch 的 CUDA 版本错位。解压后把lib目录加进环境变量export LD_LIBRARY_PATH/opt/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH pip install /opt/TensorRT-8.6.1.6/python/tensorrt-8.6.1-cp310-none-linux_x86_64.whl我用的这台机器是 TensorRT 8.6如果你下到 10.x配置方式完全一样只是 Python wheel 的文件名不同。下表是这套资源里经过验证的版本组合组件版本说明Python3.10兼容性最好Torch 和 TRT 的 wheel 都全CUDA Toolkit12.1和 PyTorch 的 cu121 对应PyTorch2.x用 torchvision 自带的 mobilenet_v2 预训练权重TensorRT8.6 或 10.x8.6 稳定10.x 支持更多新算子onnx / onnxsim最新用于导出和简化模型2.2 MobileNetV2 结构拆解倒残差块、深度可分离卷积与线性瓶颈MobileNetV2 的核心是倒残差块Inverted Residual Block和 ResNet 的残差块方向相反——先 1x1 升维再做 3x3 深度可分离卷积最后 1x1 降维降维后不跟 ReLU改叫线性瓶颈。这个设计是为了避免低维特征经过 ReLU 后被截断丢失信息。import torch import torch.nn as nn class ConvBNReLU(nn.Sequential): def __init__(self, in_planes, out_planes, kernel_size3, stride1, groups1): padding (kernel_size - 1) // 2 super().__init__( nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groupsgroups, biasFalse), nn.BatchNorm2d(out_planes), nn.ReLU6(inplaceTrue) ) class InvertedResidual(nn.Module): def __init__(self, in_planes, out_planes, stride, expand_ratio): super().__init__() hidden_dim int(round(in_planes * expand_ratio)) self.use_residual stride 1 and in_planes out_planes layers [] if expand_ratio ! 1: layers.append(ConvBNReLU(in_planes, hidden_dim, kernel_size1)) layers.extend([ ConvBNReLU(hidden_dim, hidden_dim, kernel_size3, stridestride, groupshidden_dim), nn.Conv2d(hidden_dim, out_planes, kernel_size1, biasFalse), nn.BatchNorm2d(out_planes) ]) self.conv nn.Sequential(*layers) def forward(self, x): if self.use_residual: return x self.conv(x) return self.conv(x)expand_ratio控制升维倍数MobileNetV2 里大部分块是 6 倍也就是把通道数临时放大 6 倍再压缩回来参数量没增加但每个卷积处理的信息更多。groupshidden_dim是深度可分离卷积分组卷积的写法每个通道独立卷这是 MobileNet 系列参数量远小于 VGG、ResNet 的关键。实际工程里不需要手写整个网络torchvision 已经内置了完整结构上面这段代码是用来理解结构、方便改造成自定义变体的。如果你只是做分类任务直接调 torchvision 的接口就行后文训练部分就用内置版本。2.3 数据加载与预处理归一化参数为什么是 0.485、0.456、0.406图像分类的预处理有一套固定的统计值来自 ImageNet 数据集的 RGB 三通道均值和标准差。很多人会在这里翻车——训练时用了一组归一化参数部署时忘了用同一组结果推理结果和训练时完全对不上。from torchvision import datasets, transforms from torch.utils.data import DataLoader train_tf 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]) ]) val_tf transforms.Compose([ transforms.Resize((256, 256)), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_data datasets.ImageFolder(dataset/train, transformtrain_tf) val_data datasets.ImageFolder(dataset/val, transformval_tf) train_loader DataLoader(train_data, batch_size64, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_data, batch_size64, shuffleFalse, num_workers4, pin_memoryTrue)RandomCrop(224)配合Resize((256, 256))是训练集常用的数据增强组合先把图放大到 256再随机裁 224相当于给模型提供了一定程度的平移不变性。验证集不能随机所以用CenterCrop。DataLoader里的num_workers4在 Linux 下能让数据读取和 GPU 训练并行Windows 下要注意多进程有时会报错建议改成 0 或 2。pin_memoryTrue对 GPU 训练有加速效果配合to(device, non_blockingTrue)能减少数据传输阻塞。3. 迁移学习训练用预训练权重把收敛时间压缩到三分之一3.1 数据集组织与类别映射这套资源默认的数据集目录结构是标准的ImageFolder格式训练集和验证集分开每个子文件夹代表一个类别文件夹名就是类别名。dataset/ ├── train/ │ ├── cat/ │ ├── dog/ │ └── bird/ └── val/ ├── cat/ ├── dog/ └── bird/datasets.ImageFolder会自动按文件夹名生成类别索引train_data.classes就是类别列表train_data.class_to_idx是类别到索引的映射字典。训练之前先打印出来确认一下避免类别顺序乱掉导致评估指标对不上号。如果你的数据是 CSV 标注或者 JSON 格式需要先转换成这个目录结构。数据量少的时候比如每类只有几十张Transfer Learning 依然能跑这就是热词里提到的 1-shot、5-shot 小样本场景。每类只有一张图时预训练权重还能撑住但学习率必须压到 1e-4 以下而且只能微调最后一层分类头。3.2 训练脚本迁移学习、权重衰减、标签平滑与学习率调度训练的核心是加载 ImageNet 预训练权重然后替换最后的全连接层。MobileNetV2 的classifier是一个Sequential由 Dropout 和 Linear 组成替换索引为 1 的那一层即可。import torch import torch.nn as nn import torch.optim as optim from torchvision import models num_classes len(train_data.classes) model models.mobilenet_v2(weightsmodels.MobileNet_V2_Weights.IMAGENET1K_V1) # 冻结特征提取层只训练分类头适合数据量少的场景 for param in model.features.parameters(): param.requires_grad False model.classifier[1] nn.Linear(model.last_channel, num_classes) model model.to(device) criterion nn.CrossEntropyLoss(label_smoothing0.1) optimizer optim.SGD(model.parameters(), lr0.001, momentum0.9, weight_decay5e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max50)label_smoothing0.1是标签平滑把 one-hot 的 1 和 0 变成 0.9 和 0.1/num_classes防止模型对训练集过度自信泛化会好一些。weight_decay5e-4就是 L2 正则化在 PyTorch 里的落地写法SGD 的权重衰减直接加在参数更新上能抑制过拟合。冻结model.features后优化器只更新分类头训练速度快很多。如果你的数据集每一类超过 500 张图可以解冻部分特征层把学习率设成 1e-4效果会更好。判断标准很简单训练集准确率很快到 99%、验证集卡在 80% 附近就是过拟合此时加强数据增强或者加大 weight_decay如果两边都上不去说明学习率太高或者数据有问题。余弦退火CosineAnnealingLR让学习率从初始值平滑降到接近 0比固定学习率在分类任务上平均能涨 1~2 个点。T_max50要和训练轮数一致。3.3 训练循环保存最佳模型与日志记录训练循环本身不复杂关键是每轮结束要跑验证集并且只在验证准确率创新高时保存权重这样即使后面训练发散手里也有一份最优模型。best_acc 0.0 for epoch in range(50): 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) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() acc correct / total scheduler.step() if acc best_acc: best_acc acc torch.save(model.state_dict(), mobilenetv2_best.pth) print(fepoch {epoch1} | loss {running_loss/len(train_loader):.4f} | val_acc {acc:.4f})torch.max(outputs, 1)返回每个样本得分最高的类别索引predicted labels得到布尔掩码sum().item()统计正确数。保存的mobilenetv2_best.pth是纯状态字典加载时要先建模型再load_state_dict这个文件也是后续转 ONNX 的输入。50 轮训练在单张 RTX 3060 上用 batch size 64 跑差不多 30 到 40 分钟ImageNet 预训练权重起跑收敛曲线会非常漂亮。如果发现验证集准确率一直徘徊在 60% 以下先检查数据预处理——归一化参数用错或者图片尺寸不一致是最常见的原因。4. 模型转换PyTorch 到 ONNX再到 TensorRT 的完整落盘流程4.1 导出 ONNX动态 batch 与 opset 版本选择训练完成后的模型是.pth状态字典TensorRT 不认识这个格式需要先导出成 ONNX。导出前必须把模型切到 eval 模式并且固定 batch 维度或者显式声明动态轴。import torch from torchvision import models model models.mobilenet_v2() model.classifier[1] torch.nn.Linear(model.last_channel, 3) model.load_state_dict(torch.load(mobilenetv2_best.pth, map_locationcpu)) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, mobilenetv2.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version17, do_constant_foldingTrue )dynamic_axes声明第 0 维是动态的这样导出的 ONNX 模型可以接受 batch 为 1、4、8、16 任意值。如果你的 TensorRT 版本是 8.6 之前的opset_version17可能解析不了改成 13 就行——不要在 TensorRT 装得老的情况下追求高 opset能吃进去比什么都重要。do_constant_foldingTrue会把能算的常量折叠掉减小模型体积。导出后用 onnxsim 做一次简化能去掉很多 TensorRT 不友好的多余节点pip install onnxsim onnx python -m onnxsim mobilenetv2.onnx mobilenetv2_sim.onnx简化前后的模型在推理结果上完全一致但简化后的 ONNX 转 TensorRT 成功率更高。这个步骤是血泪经验我之前跳过 onnxsim结果 TensorRT 报FAILED: Assertion failed: inputs.size() outputs.size()简化后一次通过。4.2 构建 TensorRT 引擎trtexec 与 Python API 两种方式TensorRT 引擎有两种构建方式。命令行工具 trtexec 适合快速验证和测性能Python API 适合集成到自动化流程里。先看 trtexec 的用法trtexec --onnxmobilenetv2_sim.onnx --saveEnginemobilenetv2_fp16.engine \ --fp16 \ --minShapesinput:1x3x224x224 \ --optShapesinput:8x3x224x224 \ --maxShapesinput:16x3x224x224 \ --workspace2048minShapes、optShapes、maxShapes三个参数对应动态 batch 的三种输入形状TensorRT 会按optShapes做最优 kernel 选择。--workspace2048单位是 MiB也就是给 TensorRT 最多 2GB 显存做层融合和 kernel 选择设太大会让构建时显存暴涨设太小可能触发cannot allocate memory。Python API 构建的好处是可以把构建过程封装进训练脚本训练完成直接出引擎import tensorrt as trt logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, logger) with open(mobilenetv2_sim.onnx, rb) as f: success parser.parse(f.read()) if not success: for i in range(parser.num_errors): print(parser.get_error(i)) config builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 2 * 1024 * 1024 * 1024) config.set_flag(trt.BuilderFlag.FP16) engine builder.build_serialized_network(network, config) with open(mobilenetv2_fp16.engine, wb) as f: f.write(engine)注意EXPLICIT_BATCH这个 flagONNX 导出时如果用了dynamic_axes构建网络就必须开显式 batch 模式否则 TensorRT 会认为 batch 维是固定的。build_serialized_network返回的是序列化后的引擎字节流直接写盘部署时反序列化就能用。4.3 推理验证引擎加载与精度一致性检查引擎构建完先别急着上生产跑一次推理对比 PyTorch 原始输出确认精度没有丢失。这一步在部署时是强制流程防止 FP16 转出来结果完全不对还在排查下游业务。import numpy as np import tensorrt as trt logger trt.Logger(trt.Logger.WARNING) runtime trt.Runtime(logger) with open(mobilenetv2_fp16.engine, rb) as f: engine runtime.deserialize_cuda_engine(f.read()) context engine.create_execution_context() # 动态 batch 必须显式设置输入 shape context.set_input_shape(input, (1, 3, 224, 224)) # 分配输入输出缓冲区 input_buf np.random.randn(1, 3, 224, 224).astype(np.float32)加载引擎后context.set_input_shape是动态 batch 推理的必需步骤不设会直接报INVALID_SHAPE。TensorRT 10.x 里缓冲区分配方式有变化推荐用tensorrt.tensorrt.Input的allocate_array接口但 8.6 用cuda.mem_alloc更稳。精度对比的做法是把同一张预处理后的图片分别喂给 PyTorch 模型和 TensorRT 引擎比较 softmax 输出的最大差异。差异小于 1e-2 说明转换没问题如果输出全是 NaN 或者分类结果完全不同大概率是预处理不一致看下一章的避坑指南。5. TensorRT 部署避坑指南五条实战踩坑记录5.1 ONNX 导出成功但 TensorRT 解析时报错现象trtexec 加载 ONNX 时报Unsupported operator或者Graph conflicts with provided input shape还没开始构建引擎就退出。原因opset 版本和 TensorRT 支持的算子不匹配或者 ONNX 图里残留了未折叠的冗余节点。我在 TensorRT 8.6 上试过 opset 17 的某些导出某些版本的 onnx 导出器产生的Resize节点 TRT 解析不了。解决先把 opset 降到 13 重新导出再跑一遍 onnxsim 简化。如果还报错用python -m onnxruntime.tools.convert_onnx_models_to_onnx做版本转换或者直接换 TensorRT 10.x。另外把do_constant_foldingTrue打开能消掉大部分冗余常量节点。5.2 引擎构建成功推理结果全是错的现象TensorRT 引擎能正常输出但输出的置信度分布和 PyTorch 完全对不上分类结果变成随机猜测。原因这不是 TensorRT 的问题是预处理不一致。训练时 PyTorch 侧做了ToTensorNormalize但部署侧喂给 TensorRT 的是 0~255 的原始图像或者归一化参数写错。解决部署脚本里把预处理写在 numpy 层保证和训练侧完全一致。我一般会写一个preprocess()函数输入 BGR 或 RGB 图像输出(1, 3, 224, 224)的 float32 数组里面带上 Resize、CenterCrop、Normalize然后把这个函数的输出直接作为 TensorRT 的输入。比对脚本用np.max(np.abs(trt_output - torch_output))检查差异。5.3 FP16 精度掉点明显准确率从 95% 掉到 88%现象FP16 引擎跑起来速度很快但验证集准确率明显下降集中在某几个类别上。原因MobileNetV2 的倒残差结构里有大量 BatchNorm 层FP16 下部分通道的动态范围不够激活值被截断导致特征信息丢失。这种情况通常只出现在个别 sensitive layer 上不是全模型都坏。解决用 TensorRT 的层黑名单机制把敏感层强制留在 FP32。做法是先构建一个全部 FP16 的引擎用验证集找出分类错误集中的层然后在构建配置里对这些层设置precision trt.float32。更快的做法是直接跑一遍校准给这些层加上set_precision接口for i in range(network.num_layers): layer network.get_layer(i) if layer.name in sensitive_layers: layer.precision trt.float32 layer.set_output_type(0, trt.float32)如果不想逐层排查还有一个通用方案FP16 整体掉点超过 1% 时改用 INT8 加校准数据集精度通常能回到 FP32 水平延迟只比 FP16 高一点点。5.4 动态 batch 推理时显存直接爆掉现象引擎能加载单张图推理正常但并发请求一上来就CUDA out of memory。原因构建引擎时maxShapes设成了 16workspace又给了 2GBTensorRT 会为最大 batch 预留 kernel workspace。实际业务并发根本到不了 16白占显存。解决按真实业务峰值设置optShapes和maxShapes。如果线上并发高峰是 4maxShapes设 8 就够了多出来的空间留给 TensorRT 做 kernel 调优。workspace 设置成1 301GB对 MobileNetV2 这种轻量网络已经非常充裕没必要给 2GB。显存占用和延迟是跷跷板没有绝对最优只能按场景调。5.5 PowerShell 里 conda 命令无法识别现象在 Windows PowerShell 里执行conda activate mobilenetv2终端直接报无法将“conda”项识别为 cmdlet、函数、脚本文件或可运行程序的名称。原因conda 安装时没有写入 PowerShell 的环境变量或者没有执行过conda init。这个坑在开发环境配置里出现频率极高。解决先执行conda init powershell然后关闭重开终端。如果还不行直接用 Anaconda 安装路径下的完整 python 路径C:\Users\xxx\anaconda3\envs\mobilenetv2\python.exe -m pip install torch torchvision绕过 conda activate直接调环境里的 python 解释器效果一样。这套资源里部署脚本全部用相对路径和命令行参数不会依赖特定的 shell 环境。6. 端到端性能验证用 CUDA Event 给推理延迟做个体检部署完成后的第一件事不是看 FPS而是用 CUDA Event 测端到端延迟。CUDA Event 走的是 GPU 时间线比 Python 层面的time.time()准得多尤其能过滤掉 CPU 侧数据拷贝和 Python 解释器调度的抖动。import torch import numpy as np import tensorrt as trt import pycuda.autoinit import pycuda.driver as cuda # 假设 engine 和 context 已经加载input_buf 和 output_buf 已分配 start torch.cuda.Event(enable_timingTrue) end torch.cuda.Event(enable_timingTrue) # 先预热 20 次把 TensorRT 的 kernel 加载和 CUDA context 初始化排除掉 for _ in range(20): context.execute_v2(bindings[input_buf, output_buf]) torch.cuda.synchronize() times [] for _ in range(100): start.record() context.execute_v2(bindings[input_buf, output_buf]) end.record() torch.cuda.synchronize() times.append(start.elapsed_time(end)) print(favg: {np.mean(times):.2f} ms | p95: {np.percentile(times, 95):.2f} ms)p95比平均延迟更能反映真实用户体验因为偶发的调度抖动会把平均值拉高。同一台机器上PyTorch 的 FP32 推理延迟通常在 8~12msTensorRT FP16 能压到 2~4ms吞吐提升接近 3 倍。显存占用也会跟着降因为 TRT 的 kernel 融合减少了中间张量存储。验证精度的脚本我建议做成固定的回归测试跑完 val 集把 TensorRT 输出的 top-1 准确率和 PyTorch 训练时的结果对比偏差超过 0.5% 就报警。从那以后我每次做完模型转换都强制走一遍「加载引擎 → 精度回归 → p95 延迟」三段式体检不再凭感觉上线。这套资源里的脚本已经把三段流程串好了换数据路径、换类别数就能直接用。希望帮到你。本文还有配套的精品资源点击获取