| 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=False)
- # model = shufflenet_v2_x1_0()
- model_name = "resnet50"
- model.fc = nn.Linear(int(model.fc.in_features), 2, bias=True)
- model.load_state_dict(torch.load(rf'./{model_name}.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,
- f"{model_name}.onnx",
- opset_version=13,
- export_params=True,
- do_constant_folding=True,
- ) # xxx.onnx表示.onnx文件, 这一步导出为onnx模型, 并不做任何算子融合操作。
- # 验证模型
- onnx_model = onnx.load(f"{model_name}.onnx") # 使用不同变量名
- onnx.checker.check_model(onnx_model) # 验证模型完整性
- # 使用ONNX Runtime进行推理
- ort_session = ort.InferenceSession(f"{model_name}.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]
|
