#推理加速框架：从认知到落地的核心组件（本文带OpenVINO/TensorRT入门彩蛋）

import torch
import torchvision.models as models

# 1. 加载预训练模型并切换到推理模式
model = models.resnet18(pretrained=True).eval()
# 2. 准备一个“假输入”（用来告诉ONNX模型的输入形状、数据类型等信息）
dummy_input = torch.randn(1, 3, 224, 224)  # Batch=1, 3通道RGB, 224x224尺寸
# 3. 模型转换
onnx_path = "resnet18.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    export_params=True,      # 导出模型的权重参数
    opset_version=13,        # ONNX算子集版本（尽量选12以上，兼容性好）
    do_constant_folding=True,# 常量折叠优化（提前计算好可以计算的常量）
    input_names=["input"],   # 给输入张量起个名字
    output_names=["output"], # 给输出张量起个名字
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}  # 动态Batch支持
)
print(f"PyTorch模型已成功转换为ONNX格式，保存路径：{onnx_path}")

import cv2
import numpy as np
from openvino.runtime import Core

# 1. 初始化OpenVINO推理核心
ie = Core()
# 2. 加载ONNX模型并编译到指定设备（这里用CPU，也可以用GPU、NPU等）
model = ie.read_model(model=onnx_path)
compiled_model = ie.compile_model(model=model, device_name="CPU")
# 3. 创建推理请求
infer_request = compiled_model.create_infer_request()
# 4. 准备输入数据（这里用一张随机生成的RGB图片，实际应用中换成自己的图片）
image = np.random.randn(1, 3, 224, 224).astype(np.float32)
# 5. 执行推理（同步推理，异步推理适合高吞吐量场景）
output = infer_request.infer(inputs={"input": image})
# 6. 获取输出结果（ResNet18的输出是1000个类别的概率）
output_tensor = output[compiled_model.output(0)]
predicted_class = np.argmax(output_tensor)
print(f"OpenVINO推理完成，预测类别：{predicted_class}")

import tensorrt as trt
import numpy as np

# 1. 初始化TensorRT Logger（用来输出日志信息）
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_file_path):
    """
    构建TensorRT Engine的辅助函数
    """
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, TRT_LOGGER)
    config = builder.create_builder_config()
    config.max_workspace_size = 1 << 30  # 分配1GB的工作空间给TensorRT
    # 支持FP16量化（如果GPU支持的话）
    if builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    # 解析ONNX模型
    with open(onnx_file_path, "rb") as f:
        if not parser.parse(f.read()):
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None
    # 构建并序列化TensorRT Engine
    return builder.build_engine(network, config)

# 2. 构建TensorRT Engine
engine = build_engine(onnx_path)
if engine is None:
    print("TensorRT Engine构建失败！")
    exit()
# 3. 创建推理上下文
context = engine.create_execution_context()
# 4. 准备输入输出内存（Host端和Device端）
def allocate_buffers(engine):
    inputs = []
    outputs = []
    bindings = []
    stream = torch.cuda.Stream()
    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # 分配Host端内存
        host_mem = np.zeros(size, dtype=dtype)
        # 分配Device端内存
        device_mem = torch.tensor(host_mem).cuda().contiguous()
        bindings.append(int(device_mem.data_ptr()))
        if engine.binding_is_input(binding):
            inputs.append({"host": host_mem, "device": device_mem})
        else:
            outputs.append({"host": host_mem, "device": device_mem})
    return inputs, outputs, bindings, stream

inputs, outputs, bindings, stream = allocate_buffers(engine)
# 5. 准备输入数据并拷贝到Device端
image = np.random.randn(1, 3, 224, 224).astype(np.float32).flatten()
np.copyto(inputs[0]["host"], image)
inputs[0]["device"].copy_(torch.from_numpy(inputs[0]["host"]))
# 6. 执行推理（同步推理，用torch.cuda.synchronize()等待推理完成）
context.execute_async_v2(bindings=bindings, stream_handle=stream.cuda_stream)
torch.cuda.synchronize()
# 7. 把输出结果拷贝回Host端
for output in outputs:
    output["device"].copy_(torch.from_numpy(output["host"]))
    np.copyto(output["host"], output["device"].cpu().numpy())
# 8. 获取输出结果
predicted_class = np.argmax(outputs[0]["host"])
print(f"TensorRT推理完成，预测类别：{predicted_class}")