1. 基于ResNet18的钢材表面缺陷分类模型移植

任务目标

将训练好的ResNet18钢材表面缺陷分类模型（PyTorch训练）移植到ONNX Runtime环境进行推理部署，实现单张图像的缺陷类别预测，并验证移植前后模型输出的一致性。

环境与依赖

• Python 3.8+，安装：torch、torchvision、onnx、onnxruntime、numpy、Pillow
• 预训练权重：使用torchvision自带的ResNet18预训练权重（非钢材数据微调），演示完整移植流程

核心实现

以下代码完成：模型定义与权重加载 → 导出ONNX → ONNX Runtime加载推理 → 输出一致性比对。使用一张随机生成的模拟图像作为输入，避免外部数据依赖。

import torch
import torchvision.models as models
import torch.onnx
import onnxruntime as ort
import numpy as np
from PIL import Image

# 1. 定义模型并加载预训练权重（此处用ImageNet权重演示流程）
model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
model.eval()

# 2. 构造模拟输入（3x224x224，归一化到[0,1]）
dummy_input = torch.rand(1, 3, 224, 224)

# 3. 导出ONNX
onnx_path = "resnet18_steel.onnx"
torch.onnx.export(
    model, dummy_input, onnx_path,
    input_names=["input"], output_names=["output"],
    opset_version=17,
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}
)

# 4. ONNX Runtime推理
ort_session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
ort_input = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_output = ort_session.run(None, ort_input)[0]

# 5. PyTorch推理（用于一致性比对）
with torch.no_grad():
    torch_output = model(dummy_input).numpy()

# 6. 验证：比较输出差异
max_diff = np.max(np.abs(ort_output - torch_output))
print(f"ONNX与PyTorch输出最大误差: {max_diff:.6f}")

# 7. 模拟单张图像推理（实际使用时替换为真实图像路径）
def predict_image(image_path):
    img = Image.open(image_path).convert("RGB").resize((224, 224))
    img_array = np.array(img).astype(np.float32) / 255.0
    img_tensor = torch.from_numpy(img_array).permute(2, 0, 1).unsqueeze(0)
    # 使用ImageNet标准化参数
    mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
    std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
    img_tensor = (img_tensor - mean) / std
    ort_input = {ort_session.get_inputs()[0].name: img_tensor.numpy()}
    result = ort_session.run(None, ort_input)[0]
    pred_class = int(np.argmax(result[0]))
    return pred_class, result[0]

# 示例：使用随机噪声图验证推理流程可运行
sample_img = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))
sample_img.save("sample_test.png")
pred, scores = predict_image("sample_test.png")
print(f"预测类别索引: {pred}, 置信度分数shape: {scores.shape}")

运行与验证

1. 将上述代码保存为migrate_resnet18.py，直接运行：python migrate_resnet18.py
2. 观察输出：max_diff应小于1e-4（浮点精度范围内），证明ONNX Runtime与PyTorch推理结果一致
3. 程序会生成resnet18_steel.onnx模型文件和sample_test.png测试图，并打印预测类别索引

通过条件：程序无报错退出，且打印的最大误差小于1e-4。若误差过大，检查onnxruntime版本与opset兼容性。

常见问题

• 导出失败：升级torch版本至2.0+，或降低opset_version至11
• 推理速度慢：确认使用CPUExecutionProvider，或改用CUDAExecutionProvider（需GPU环境）
• 真实钢材数据替换：将predict_image中的图像预处理改为训练时的标准化参数，并将模型权重替换为微调后的.pth文件