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]