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, squeezenet1_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 = shufflenet_v2_x2_0()
    # model = squeezenet1_0()
    model_name = "shufflenet"
    if model_name == "squeezenet":
        # 获取SqueezeNet的最后一个卷积层的输入通道数
        final_conv_in_channels = model.classifier[1].in_channels
        # 替换classifier为新的Sequential，将输出改为2类
        model.classifier = nn.Sequential(
            nn.Dropout(p=0.5),
            nn.Conv2d(final_conv_in_channels, 2, kernel_size=(1, 1)),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool2d((1, 1))
        )
    if model_name == "shufflenet":
        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]