### 将YOLOv8模型转换为TensorRT `.engine` 文件以加速推理 要将YOLOv8模型转换为TensorRT的`.engine`文件，需遵循以下步骤。这些步骤涵盖了从原始模型导出到最终部署的过程。 --- #### 1. 导出YOLOv8模型为ONNX格式 首先，需要将YOLOv8模型从PyTorch的`.pt`格式导出为中间表示（Intermediate Representation, IR），即ONNX格式。这是后续转换至TensorRT的关键步骤之一[^3]。 ```python from ultralytics import YOLO # 加载YOLOv8模型 model = YOLO("yolov8n.pt") # 导出模型为 ONNX 格式 success = model.export(format="onnx", simplify=True) ``` 上述代码会生成一个名为`yolov8n.onnx`的文件，该文件将是下一步输入给TensorRT的基础[^1]。 --- #### 2. 安装TensorRT及其Python绑定工具 确保已安装NVIDIA TensorRT库以及相应的Python接口。如果尚未安装，可通过官方文档指导完成配置[^1]。 ```bash pip install nvidia-tensorrt ``` --- #### 3. 编写脚本将ONNX模型转化为TensorRT引擎 编写一段Python脚本来执行具体的转化工作。下面是一个典型的例子： ```python import tensorrt as trt import os def build_engine(onnx_file_path, engine_file_path): TRT_LOGGER = trt.Logger(trt.Logger.WARNING) with trt.Builder(TRT_LOGGER) as builder, \ builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \ trt.OnnxParser(network, TRT_LOGGER) as parser: # 设置最大批大小和工作空间 builder.max_batch_size = 1 config = builder.create_builder_config() config.max_workspace_size = 1 << 30 # 1GB 工作内存分配 # 解析 ONNX 模型 with open(onnx_file_path, 'rb') as model: if not parser.parse(model.read()): print('Failed to parse the ONNX file.') for error in range(parser.num_errors): print(parser.get_error(error)) return None # 构建 TensorRT 引擎 serialized_engine = builder.build_serialized_network(network, config) # 保存序列化的引擎到指定路径 with open(engine_file_path, "wb") as f: f.write(serialized_engine) return serialized_engine # 输入/输出文件定义 onnx_model_path = "./yolov8n.onnx" tensorrt_engine_path = "./yolov8n.engine" # 转化过程启动 build_engine(onnx_model_path, tensorrt_engine_path) print(f"Engine saved at {tensorrt_engine_path}") ``` 这段代码实现了从ONNX文件创建并存储TensorRT引擎的功能。 --- #### 4. 测试生成的TensorRT引擎 一旦成功构建了`.engine`文件，就可以通过加载这个文件来进行推理测试。如下所示： ```python import cv2 import numpy as np import pycuda.driver as cuda import pycuda.autoinit import tensorrt as trt class HostDeviceMem(object): """简单封装主机端与设备端内存""" def __init__(self, host_mem, device_mem): self.host = host_mem self.device = device_mem def __str__(self): return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device) def __repr__(self): return self.__str__() def allocate_buffers(engine): inputs = [] outputs = [] bindings = [] stream = 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)) # Allocate host and device buffers. host_mem = cuda.pagelocked_empty(size, dtype) device_mem = cuda.mem_alloc(host_mem.nbytes) # Append the device buffer to device bindings. bindings.append(int(device_mem)) # Append to input or output list. if engine.binding_is_input(binding): inputs.append(HostDeviceMem(host_mem, device_mem)) else: outputs.append(HostDeviceMem(host_mem, device_mem)) return inputs, outputs, bindings, stream def do_inference(context, bindings, inputs, outputs, stream, batch_size=1): [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle) [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] stream.synchronize() return [out.host for out in outputs] # 初始化 TensorRT 运行时环境 TRT_LOGGER = trt.Logger(trt.Logger.WARNING) with open("./yolov8n.engine", "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: engine = runtime.deserialize_cuda_engine(f.read()) context = engine.create_execution_context() inputs, outputs, bindings, stream = allocate_buffers(engine) # 准备输入数据 (假设是一张图片) image = cv2.imread('./bus.jpg') input_data = preprocess_image(image) # 自定义函数，负责图像预处理成网络所需形状 np.copyto(inputs[0].host, input_data.ravel().astype(np.float32)) # 推理阶段 output = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)[0] postprocess(output) # 后续处理结果 ``` 此处展示了如何加载之前生成的`.engine`文件，并对其进行基本推断操作[^1]。 --- #### 总结 整个流程包括四个主要部分：先是从YOLOv8模型导出ONNX格式；其次是借助TensorRT SDK把ONNX转译成高效运行于GPU上的`.engine`文件；再次是对新产生的引擎做初步验证性试验；最后则是实际项目集成环节[^1]^。 ---