#推理加速框架：ONNX Runtime、TensorRT、OpenVINO详解

import torch
import torch.onnx
import onnx

class ExampleCNN(torch.nn.Module):
    """简单的分类CNN，仅作示例"""
    def __init__(self):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(3, 64, 3, padding=1)
        self.bn1 = torch.nn.BatchNorm2d(64)
        self.relu = torch.nn.ReLU()
        self.pool = torch.nn.AdaptiveAvgPool2d((1, 1))
        self.fc = torch.nn.Linear(64, 10)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.pool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

# 1. 初始化模型，切换到评估模式（冻结BN、Dropout）
model = ExampleCNN().eval()

# 2. 构建一个虚拟输入，形状必须和真实输入一模一样
dummy_input = torch.randn(1, 3, 224, 224)  # 单张图，3通道，224x224

# 3. 导出 ONNX
torch.onnx.export(
    model,
    dummy_input,
    "example_cnn.onnx",
    export_params=True,          # 把训练好的权重一同保存
    opset_version=14,            # 算子集版本，14以上支持更多优化特性
    do_constant_folding=True,    # 折叠常量运算，例如把BN的归一化提前算掉
    input_names=["input"],       # 给输入节点起个好名字
    output_names=["output"],     # 给输出节点起个好名字
    dynamic_axes={               # 允许推理时改变 batch 大小
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)

# 4. 用onnx库验证模型是否合法
onnx_model = onnx.load("example_cnn.onnx")
onnx.checker.check_model(onnx_model)
print("✅ ONNX模型验证成功！")

import onnxruntime as ort
import numpy as np

# 1. 创建推理会话，让Runtime自己挑选最佳硬件（先CUDA后CPU）
session = ort.InferenceSession(
    "example_cnn.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
)

# 2. 查看模型期望的输入输出名字和形状
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
input_shape = session.get_inputs()[0].shape  # 如 (None, 3, 224, 224)，None代表动态batch

# 3. 准备数据（注意：onnxruntime默认吃float32，要和模型设定一致）
input_data = np.random.randn(4, 3, 224, 224).astype(np.float32)  # 一次推理4张图

# 4. 执行推理
results = session.run([output_name], {input_name: input_data})
output = results[0]
print(f"✅ 推理完成，输出形状：{output.shape}")

def get_optimized_session(model_path):
    # 开启所有内置图优化
    sess_options = ort.SessionOptions()
    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    # CPU场景：多线程绑定，通常设置为物理核数-1
    sess_options.intra_op_num_threads = 4
    sess_options.inter_op_num_threads = 1
    # 内存优化选项，开启后能重用临时张量，降低内存占用
    sess_options.enable_mem_pattern = True
    sess_options.enable_mem_reuse = True

    return ort.InferenceSession(
        model_path,
        sess_options=sess_options,
        providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
    )