TensorRT 似乎缺乏常用功能

Question

我最近遇到了一个名为tensorRT的神奇工具，但由于我的笔记本电脑上没有NVIDIA GPU，所以我决定使用Google Collab来尝试这项技术。

我使用简单的 pip 命令来安装必要的库，包括用于 CUDA 管理的库

pip install nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com
pip install pycuda

安装后，一切似乎都可以使用了。然而，事实证明，一些常用的方法根本不存在。

当我尝试通过

创建

tensorRT Engine

时

builder = trt.Builder(trt.Logger(trt.Logger.INFO))
network = builder.create_network(batch_size)
engine = builder.build_cuda_engine(network)

它抛出异常，

'tensorrt.tensorrt.Builder' has no attribute 'build_cuda_engine'

，尽管事实上它应该如此。

我是否错过了一些重要的安装，或者我只是使用了一些已弃用的版本？

Answer 1

TensorRT 确实是一个相当不错的推理工具。一开始使用起来很棘手，但很快就会变得合乎逻辑。按照 github 上提供的 python 示例进行操作here。

为了解决您的特定问题，即以编程方式构建 TensorRT 引擎，请遵循以下结构：

explicit_batch = 1 << (int)(
    trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

with trt.Builder(self.TRT_LOGGER) as builder, \
    builder.create_network(explicit_batch) as network, \
        trt.OnnxParser(network, self.TRT_LOGGER) as parser:

    with open("model.onnx", 'rb') as model:
        parser.parse(model.read())

    builder.max_workspace_size = 1 << 30
    config = builder.create_builder_config()
    config.max_workspace_size = 1 << max_workspace_size
    config.set_flag(trt.BuilderFlag.GPU_FALLBACK)
    config.set_flag(trt.BuilderFlag.FP16)


    engine = builder.build_engine(network, config)
    with open("result.engine", "wb") as f:
        f.write(engine.serialize())

这是一个非常基本的代码片段，应该可以解决您当前的问题。遵循示例并遵循此结构应该可以解决您的问题。

这个 Yolov7 github 存储库 还支持 TensorRT，并具有如何导出模型的完整实现以及关于如何推断模型的 google colab。

干杯！

Answer 2

尝试使用 YOLOv8 使用 TensorRT 导出模型时也遇到此错误

model = YOLO("./best.pt")
model.export(format="TensorRT")

我使用了 TensorRT 10 和 CUDA 12.1。
修复方法是卸载 CUDA 和 TensorRT 并下载

CUDA 11.8

TensorRT 8