幽冥大陆(八十三)Python 水果识别PTH 转 ONNX 脚本

容你代码中的输入尺寸（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。

阿雪技术观

在科技发展浪潮中，我们不妨积极投身技术共享。不满足于做受益者，更要主动担当贡献者。无论是分享代码、撰写技术博客，还是参与开源项目维护改进，每一个微小举动都可能蕴含推动技术进步的巨大能量。东方仙盟是汇聚力量的天地，我们携手在此探索硅基生命，为科技进步添砖加瓦。

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