容你代码中的输入尺寸(100x100)和 CPU 环境:
python
运行
import torch import torch.nn as nn from torchvision import models from torchvision.models import MobileNet_V2_Weights # ==================== 配置项(需和训练脚本一致) ==================== PTH_MODEL_PATH = "wlzcfruit_mobilenetv2.pth" # 训练好的pth路径 ONNX_MODEL_PATH = "wlzcfruit_mobilenetv2.onnx" # 输出onnx路径 NUM_CLASSES = 208 # 你的水果类别数(和训练时一致) INPUT_SIZE = (100, 100) # 输入尺寸(和训练时的Resize一致) DEVICE = torch.device("cpu") # 保持和训练一致的CPU环境 # =================================================================== # 1. 重建和训练时完全一致的模型结构 def build_model(num_classes): # 重建MobileNetV2(和训练脚本结构一致,消除pretrained警告) model = models.mobilenet_v2(weights=MobileNet_V2_Weights.IMAGENET1K_V1) # 修改分类头(和训练脚本一致) model.classifier[1] = nn.Linear(model.last_channel, num_classes) return model # 2. 加载pth权重并设置为推理模式 model = build_model(NUM_CLASSES) # 加载权重(CPU环境无需map_location) model.load_state_dict(torch.load(PTH_MODEL_PATH, map_location=DEVICE)) model.eval() # 必须设为推理模式,禁用Dropout/BatchNorm训练行为 # 3. 构造示例输入(匹配模型输入维度:batch_size=1, 3通道, 100x100) dummy_input = torch.randn(1, 3, INPUT_SIZE[0], INPUT_SIZE[1], device=DEVICE) # 4. 转换PTH到ONNX(核心步骤) torch.onnx.export( model, # 加载好权重的模型 dummy_input, # 示例输入(用于推断维度) ONNX_MODEL_PATH, # 输出ONNX路径 opset_version=12, # ONNX算子版本(兼容大部分推理框架) input_names=["input"], # 输入节点名称(方便后续推理调用) output_names=["output"], # 输出节点名称(分类结果) dynamic_axes={ # 支持动态批次(可选,推荐) "input": {0: "batch_size"}, "output": {0: "batch_size"} }, verbose=False # 关闭详细日志(如需调试可设为True) ) # 5. 验证ONNX模型是否有效(可选,推荐) try: import onnx # 加载并检查ONNX模型完整性 onnx_model = onnx.load(ONNX_MODEL_PATH) onnx.checker.check_model(onnx_model) print(f"✅ ONNX模型转换成功!路径:{ONNX_MODEL_PATH}") print(f"📌 模型输入维度:1x3x{INPUT_SIZE[0]}x{INPUT_SIZE[1]}(支持动态批次)") print(f"📌 模型输出维度:1x{NUM_CLASSES}(对应{NUM_CLASSES}类水果)") except ImportError: print("⚠️ 未安装onnx库,跳过模型验证(可执行 pip install onnx 安装)") print(f"✅ ONNX模型已生成:{ONNX_MODEL_PATH}") except Exception as e: print(f"❌ ONNX模型验证失败:{str(e)}")关键注意事项(必看)
参数一致性:
NUM_CLASSES必须和训练时的类别数一致(你的日志中是 208 类);INPUT_SIZE必须和训练脚本中Resize((100, 100))一致,否则推理会报错;- 若训练时修改过输入尺寸(如 64x64),需同步修改此处的
INPUT_SIZE。
环境依赖:
- 安装依赖:
pip install onnx(用于验证模型,可选但推荐); - PyTorch 版本建议和训练时一致,避免算子不兼容。
- 安装依赖:
动态批次说明:
- 脚本中
dynamic_axes配置支持动态批次(如批量推理时用 batch_size=16/32); - 若只需单张图片推理,可删除
dynamic_axes参数,模型体积会略小。
- 脚本中
推理适配:
- ONNX 模型输入需和训练时的预处理一致(若后续加回
Normalize,推理时需对输入图片执行相同归一化); - 输入图片格式要求:RGB 通道、张量形状为
(batch_size, 3, 100, 100)、数据类型为 float32。
- ONNX 模型输入需和训练时的预处理一致(若后续加回
阿雪技术观
在科技发展浪潮中,我们不妨积极投身技术共享。不满足于做受益者,更要主动担当贡献者。无论是分享代码、撰写技术博客,还是参与开源项目维护改进,每一个微小举动都可能蕴含推动技术进步的巨大能量。东方仙盟是汇聚力量的天地,我们携手在此探索硅基生命,为科技进步添砖加瓦。
Hey folks, in this wild tech - driven world, why not dive headfirst into the whole tech - sharing scene? Don't just be the one reaping all the benefits; step up and be a contributor too. Whether you're tossing out your code snippets, hammering out some tech blogs, or getting your hands dirty with maintaining and sprucing up open - source projects, every little thing you do might just end up being a massive force that pushes tech forward. And guess what? The Eastern FairyAlliance is this awesome place where we all come together. We're gonna team up and explore the whole silicon - based life thing, and in the process, we'll be fueling the growth of technology