#推理加速框架：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. 初始化模型+评估模式
model = ExampleCNN().eval()
# 2. 构建虚拟输入（形状必须匹配真实输入）
dummy_input = torch.randn(1, 3, 224, 224)  # batch_size=1, 3通道, 224x224

# 3. 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "example_cnn.onnx",
    export_params=True,          # 导出训练好的权重
    opset_version=14,            # 尽量选最新稳定的ONNX opset（14+支持更多优化）
    do_constant_folding=True,    # 自动折叠常量运算（如BN的预计算）
    input_names=["input"],       # 自定义输入节点名（方便推理时绑定）
    output_names=["output"],     # 自定义输出节点名
    dynamic_axes={                # 动态轴设置（允许推理时变batch_size）
        "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. 创建推理会话+指定硬件优先级（先试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) 动态batch

# 3. 准备真实输入数据（注意数据类型！ONNX Runtime默认用float32）
input_data = np.random.randn(4, 3, 224, 224).astype(np.float32)  # batch_size=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
    sess_options.intra_op_num_threads = 4  # 单算子内部并行线程（CPU场景设为CPU核数-1）
    sess_options.inter_op_num_threads = 1  # 多算子并行线程（无分支图设为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"]
    )