y用export生成yolov5s.engine

使用 export 命令生成 YOLOv5s 的 TensorRT 引擎文件
为了将 YOLOv5s 模型转换为 TensorRT 引擎文件 (.engine)，通常需要遵循一系列特定步骤来准备环境并执行必要的构建过程。然而，在提到的参考资料中并未直接提及通过 export 命令完成此操作的方法[^1]。
一般情况下，创建 .engine 文件涉及以下几个方面的工作：

准备工作：确保已经安装了 PyCUDA 和 TensorRT，并验证其正确安装[^5]。

编译项目：对于基于 C++ 实现的 TensorRT 推理程序（如 tensorrtx/yolov5），需先建立对应的构建目录并运行 CMake 及 Make 工具来进行编译工作[^2]。

权重文件与源码调整：确认目标路径存在相应的权重量化参数文件（.wts），并且适当更新了相关联的 C++ 或头文件中的配置项以匹配所使用的具体版本或变体（比如 yolov5s）[^4]。

生成引擎文件：最终利用预处理脚本或者可执行文件指定输入输出格式以及架构选项来生成所需的 .engine 文件。例如，可以通过如下命令实现：
sudo ./yolov5 -s ../yolov5s.wts yolov5s.engine s
登录后复制

尽管上述方法适用于某些场景下的 .engine 文件生成流程，但对于希望采用更简便的方式——即使用 Python 脚本内的 export 函数或其他高级 API 来导出模型的情况，则可能依赖于官方支持库所提供的功能或是第三方工具的支持。如果考虑使用 Python 中的相关函数进行模型转换，建议查阅最新版 TensorRT 文档获取有关如何调用此类接口的具体指导。
from models.experimental import attempt_load
import torch
from utils.general import non_max_suppression
from pathlib import Path
import onnxruntime as ort
import numpy as np

def export_to_tensorrt(model_path='weights/best.pt', engine_file_name="yolov5s.trt", batch_size=1, half=False):
device = 'cuda' if torch.cuda.is_available() else 'cpu'

model = attempt_load(weights=model_path, map_location=device)

input_names = ["images"]
output_names = ['output']

dynamic_axes = {'images': {0: 'batch'}, 'output': {0: 'batch'}}

dummy_input = torch.randn(batch_size, 3, 640, 640).to(device)

with torch.no_grad():
outputs = model(dummy_input)[0]

torch.onnx.export(
model,
dummy_input,
f"{Path(engine_file_name).stem}.onnx",
opset_version=11,
do_constant_folding=True,
input_names=input_names,
output_names=output_names,
dynamic_axes=dynamic_axes
)

sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

session = ort.InferenceSession(f"{Path(engine_file_name).stem}.onnx", providers=['TensorrtExecutionProvider'], sess_options=sess_options)

inputs = {
"images": dummy_input.cpu().numpy(),
}

_ = session.run(None, inputs)

if __name__ == "__main__":
export_to_tensorrt()
登录后复制
该代码片段展示了另一种途径，即将 PyTorch 训练好的 YOLOv5s 模型转成 ONNX 格式后再借助 ONNX Runtime 进一步优化至 TensorRT 后端。这提供了一种不同于传统 C++ 编程方式的新思路。