import torch

import torchvision.models as models

import onnx

import onnxruntime

# 加载 PyTorch 模型

model = models.resnet18(weights=True)

model.eval()

# 定义输入和输出张量的名称和形状

input_names = ["input"]

output_names = ["output"]

batch_size = 1

input_shape = (batch_size, 3, 224, 224)

output_shape = (batch_size, 1000)

# 将 PyTorch 模型转换为 ONNX 格式

torch.onnx.export(

    model,  # 要转换的 PyTorch 模型

    torch.randn(input_shape),  # 模型输入的随机张量

    "resnet18.onnx",  # 保存的 ONNX 模型的文件名

    input_names=input_names,  # 输入张量的名称

    output_names=output_names,  # 输出张量的名称

    dynamic_axes={input_names[0]: {0: "batch_size"}, output_names[0]: {0: "batch_size"}}  # 动态轴，即输入和输出张量可以具有不同的批次大小

)

# 加载 ONNX 模型

onnx_model = onnx.load("resnet18.onnx")

onnx_model_graph = onnx_model.graph

onnx_session = onnxruntime.InferenceSession(onnx_model.SerializeToString())

# 使用随机张量测试 ONNX 模型

x = torch.randn(input_shape).numpy()

onnx_output = onnx_session.run(output_names, {input_names[0]: x})[0]

print(f"PyTorch output: {model(torch.from_numpy(x)).detach().numpy()[0, :5]}")

print(f"ONNX output: {onnx_output[0, :5]}")

import tensorrt as trt

import numpy as np

# 加载 ONNX 模型

onnx_file_path = "resnet18.onnx"

onnx_model = trt.OnnxParser.create_network(

    parser=trt.OnnxParser(network=trt.Network()),  # 创建一个 ONNX 解析器

    flags=0  # 解析器标志，0 表示默认值

)

with open(onnx_file_path, "rb") as model_file:  # 打开 ONNX 模型文件

    onnx_model.deserialize(model_file.read())  # 反序列化 ONNX 模型

# 创建 TensorRT 引擎

trt_logger = trt.Logger(trt.Logger.WARNING)  # 创建一个 TensorRT 日志记录器

trt_builder = trt.Builder(trt_logger)  # 创建一个 TensorRT 构建器

trt_builder.max_batch_size = 1  # 设置最大批次大小

trt_builder.max_workspace_size = 1 << 30  # 设置最大工作空间大小

trt_builder.fp16_mode = True  # 启用 FP16 模式

trt_builder.int8_mode = False  # 禁用 INT8 模式

trt_engine = trt_builder.build_cuda_engine(onnx_model)  # 构建 TensorRT 引擎

# 保存 TensorRT 引擎

trt_engine_path = "resnet18.engine"

trt_engine.serialize().tofile(trt_engine_path)  # 将 TensorRT 引擎序列化为字节流并保存到文件中

import tensorrt as trt

import pycuda.driver as cuda

import pycuda.autoinit

import numpy as np

# 加载 TensorRT 引擎

trt_engine_path = "resnet18.trt"

trt_engine = trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine(trt_engine_path)

# 创建执行上下文

trt_context = trt_engine.create_execution_context()

# 定义输入和输出张量的形状

input_shape = (1, 3, 224, 224)

output_shape = (1, 1000)

# 分配内存

input_memory = cuda.mem_alloc(trt_engine.get_binding_shape(0)[0] * trt.volume(input_shape))

output_memory = cuda.mem_alloc(trt_engine.get_binding_shape(1)[0] * trt.volume(output_shape))

# 准备输入数据

input_data = np.random.randn(*input_shape).astype(np.float32)

# 将输入数据复制到 GPU 内存中

cuda.memcpy_htod(input_memory, input_data)

# 执行推理

trt_context.execute_v2(bindings=[int(input_memory), int(output_memory)])

# 将输出数据从 GPU 内存中复制到主机内存中

output_data = np.empty(output_shape, dtype=np.float32)

cuda.memcpy_dtoh(output_data, output_memory)

# 打印输出结果

print(output_data)