幽冥大陆(八十四)Python 水果识别PTH 转 ONNX 脚本新 —东方仙盟练气期

的 PTH 模型实际训练的是131 类 水果，但转换脚本中手动设置了NUM_CLASSES = 208，导致分类头的权重维度不匹配（131≠208），这是典型的「模型结构和权重维度不一致」问题。

修正后的完整 PTH 转 ONNX 脚本（适配 131 类）

python

运行

复制代码

import torch
import torch.nn as nn
from torchvision import models
from torchvision.models import MobileNet_V2_Weights

# ==================== 配置项（修正类别数为131） ====================
PTH_MODEL_PATH = "wlzcfruit_mobilenetv2.pth"  # 训练好的pth路径
ONNX_MODEL_PATH = "wlzcfruit_mobilenetv2.onnx"  # 输出onnx路径
NUM_CLASSES = 131  # ✅ 修正为实际训练的131类（不是208）
INPUT_SIZE = (100, 100)  # 和训练脚本的Resize一致
DEVICE = torch.device("cpu")
# ===================================================================

# 1. 重建和训练时完全一致的模型结构
def build_model(num_classes):
    # 重建MobileNetV2（和训练脚本结构完全一致）
    model = models.mobilenet_v2(pretrained=True)  # 保持和训练脚本一致的写法
    # 修改分类头（维度匹配131类）
    model.classifier[1] = nn.Linear(model.last_channel, num_classes)
    return model

# 2. 加载pth权重并设置为推理模式
model = build_model(NUM_CLASSES)
# 加载权重（CPU环境）
model.load_state_dict(torch.load(PTH_MODEL_PATH, map_location=DEVICE))
model.eval()  # 必须设为推理模式

# 3. 构造示例输入（1,3,100,100）
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,
    opset_version=12,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    },
    verbose=False
)

# 5. 验证ONNX模型
try:
    import 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"❌ 验证失败：{str(e)}")

额外优化：自动匹配类别数（避免手动填错）

如果担心后续类别数变化，可从训练集目录自动读取类别数，彻底避免维度不匹配：

python

运行

复制代码

import os
# 自动读取训练集类别数（和训练脚本逻辑一致）
TRAIN_DIR = "./wlzcfruit/Training"
NUM_CLASSES = len([d for d in os.listdir(TRAIN_DIR) if os.path.isdir(os.path.join(TRAIN_DIR, d))])
print(f"🔍 自动检测到训练集类别数：{NUM_CLASSES}")

执行步骤

替换脚本中NUM_CLASSES = 208为131（或用上面的自动检测代码）；
确保PTH_MODEL_PATH指向正确的wlzcfruit_mobilenetv2.pth；
重新运行脚本：python wlzc_fruit_pth2onnx.py。

阿雪技术观

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

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