| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647 |
- import torch
- import onnx
- import torch.onnx
- import onnxruntime as ort
- import numpy as np
- from torchvision.models import resnet50, shufflenet_v2_x1_0, shufflenet_v2_x2_0
- from torch import nn
- from simple_model import SimpleModel
- if __name__ == '__main__':
- # 载入模型框架
- # model = SimpleModel()
- # model = resnet50(pretrained=True)
- model = shufflenet_v2_x1_0()
- model.fc = nn.Linear(int(model.fc.in_features), 2, bias=False)
- model.load_state_dict(torch.load(r'./shufflenet.pth')) # xxx.pth表示.pth文件, 这一步载入模型权重
- print("加载模型成功")
- model.eval() # 设置模型为推理模式
- example_input = torch.randn(1, 3, 256, 256) # [1,3,224,224]分别对应[B,C,H,W]
- # print(model)
- torch.onnx.export(model,
- example_input,
- "shufflenet.onnx",
- opset_version=13,
- export_params=True,
- do_constant_folding=True,
- ) # xxx.onnx表示.onnx文件, 这一步导出为onnx模型, 并不做任何算子融合操作。
- # 验证模型
- onnx_model = onnx.load("shufflenet.onnx") # 使用不同变量名
- onnx.checker.check_model(onnx_model) # 验证模型完整性
- # 使用ONNX Runtime进行推理
- ort_session = ort.InferenceSession("shufflenet.onnx")
- ort_inputs = {ort_session.get_inputs()[0].name: example_input.detach().numpy()}
- ort_outs = ort_session.run(None, ort_inputs)
- # 与PyTorch原始输出对比
- with torch.no_grad():
- torch_out = model(example_input)
- # 检查最大误差
- print("输出差异最大为:", np.max(np.abs(torch_out.numpy() - ort_outs[0])))
- #mean_mlir = [0.485×255, 0.456×255, 0.406×255] = [123.675, 116.28, 103.53]
- #scale_mlir = [0.229*255, 0.224*255, 0.225*255] = [58.395, 57.12, 57.375]
|
