nvidia - Tensorrt python API 设置批量大小

Question

我想在构建 TensorRT 引擎时设置批量大小。

我正在使用 TensorRT 7 和 python API。我正在将 ResNet50 模型转换为 onnx 格式。

Answer 1

我无法在 python API 中做到这一点。但是，我使用了trtexectensorrt 默认自带的工具。该工具将 onnx 模型转换为 tensorrt 引擎。

要转换模型，请使用以下命令：

trtexec --explicitBatch --onnx=model.onnx --minShapes=input:min_batchxsample_size --maxShapes=input:max_batchxsample_size --optShapes=input:opt_batchxsample_size --saveEngine=model.engine

有关trtexec使用此 GitHub页面的更多信息。