fix(demo): revert #1786 and update trt doc (#1792)

Megvii-BaseDetection · Jul 30, 2024 · f00a798 · f00a798
1 parent 58c2dd0
commit f00a798
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/demo/TensorRT/cpp/README.md b/demo/TensorRT/cpp/README.md
@@ -46,3 +46,13 @@ or
 ./yolox <path/to/your/engine_file> -i <path/to/image>
 ```
 
+NOTE: for `trtexec` users, modify `INPUT_BLOB_NAME` and `OUTPUT_BLOB_NAME` as the following code.
+```
+const char* INPUT_BLOB_NAME = "images";
+const char* OUTPUT_BLOB_NAME = "output";
+```
+
+Here is the command to convert the small onnx model to tensorrt engine file:
+```
+trtexec --onnx=yolox_s.onnx --saveEngine=yolox_s.trt
+```
diff --git a/demo/TensorRT/cpp/yolox.cpp b/demo/TensorRT/cpp/yolox.cpp
@@ -31,8 +31,8 @@ using namespace nvinfer1;
 static const int INPUT_W = 640;
 static const int INPUT_H = 640;
 static const int NUM_CLASSES = 80;
-const char* INPUT_BLOB_NAME = "images";
-const char* OUTPUT_BLOB_NAME = "output";
+const char* INPUT_BLOB_NAME = "input_0";
+const char* OUTPUT_BLOB_NAME = "output_0";
 static Logger gLogger;
 
 cv::Mat static_resize(cv::Mat& img) {