import torch
import torch.onnx
from torchvision.models import resnet50, ResNet50_Weights
from torch import nn

if __name__ == '__main__':
    input = torch.randn(1, 3, 256, 256)          # [1,3,224,224]分别对应[B,C,H,W]
    # 载入模型框架
    model = resnet50()
    # model.fc = nn.Sequential(
    #     nn.Linear(int(model.fc.in_features), int(model.fc.in_features) // 2, bias=True),
    #     nn.ReLU(inplace=True),
    #     nn.Dropout(0.5),
    #     nn.Linear(int(model.fc.in_features) // 2, 2, bias=False)
    # )

    # model.load_state_dict(torch.load("resnet50_best_model_acc.pth")) # xxx.pth表示.pth文件, 这一步载入模型权重
    model.load_state_dict(torch.load(r'D:\code\water_turbidity_det\resnet50-11ad3fa6.pth')) # xxx.pth表示.pth文件, 这一步载入模型权重
    model.eval()                                 # 设置模型为推理模式
   # print(model)
   # model = torch.jit.script(model)  # 先转换为TorchScript
    torch.onnx.export(model,
                      input,
                      "resnet50_best_model_acc.onnx",
                      training=torch.onnx.TrainingMode.EVAL,
                      opset_version=18,
                      export_params=True,
                      do_constant_folding=True,
                      input_names=['input'],
                      output_names=['output']
                      )  # xxx.onnx表示.onnx文件, 这一步导出为onnx模型, 并不做任何算子融合操作。

    # 验证模型
    import onnx
    model = onnx.load("resnet50_best_model_acc.onnx")
    onnx.checker.check_model(model)  # 验证模型完整性

    #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]