#推理加速框架：从认知到落地的核心组件（含 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, 224×224

# 3. 导出 ONNX 模型
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"模型成功转换为 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. 准备输入数据（这里用随机生成的数据，实际使用时替换为真实图片）
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
import torch   # 仅用于 GPU 内存交互

# 1. 初始化 TensorRT Logger
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_file_path):
    """
    从 ONNX 构建 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   # 1 GB 工作空间

    # 如果 GPU 支持 FP16，则启用 FP16 优化
    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

    # 构建并序列化 Engine
    return builder.build_engine(network, config)

# 2. 构建 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. 准备输入数据并拷贝到 GPU
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. 执行异步推理，然后等待完成
context.execute_async_v2(bindings=bindings, stream_handle=stream.cuda_stream)
torch.cuda.synchronize()

# 7. 将输出从 GPU 拷贝回 CPU
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}")