devbox
小小林熬夜学编程
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

（有待完善）Ubuntu20.04安装部署TensorRT_ubuntu20.04安装tensorrt

作者：小小林熬夜学编程 | 2024-04-19 03:22:30

赞

踩

ubuntu20.04安装tensorrt

、下载解压

下载网址：https://developer.nvidia.com/nvidia-tensorrt-download

下载完拉到主目录下

tar -zxvf TensorRT-8.2.1.8.Linux.x86_64-gnu.cuda-11.4.cudnn8.2.tar.gz

二、添加环境变量

gedit ~/.bashrc

文末添加：

export LD_LIBRARY_PATH=/home/wave/TensorRT-8.2.1.8/lib:$LD_LIBRARY_PATH

刷新环境变量

source ~/.bashrc

三、下载源码

网址：https://gitcode.net/mirrors/wang-xinyu/tensorrtx?utm_source=csdn_github_accelerator

注意##下载的版本按照你安装的yolov5版本

下载yolov5对应的模型网址：https://github.com/ultralytics/yolov5/releases

上一篇博客你安装的yolov5-v5.0的话，上面的源码和模型也要对应版本。

都下载完后进入yolov5的目录

将里面的gen_wts.py文件复制到你yolov5的目录下

在此目录下打开终端生成.wts文件

python3 gen_wts.py --w yolov5s.pt

生成完毕后回到这个目录

mkdir build && cd build
cmake ..
make

然后把刚刚生成的.wts文件拉到这个build文件下

sudo ./yolov5 -s yolov5s.wts yolov5s.engine s
sudo ./yolov5 -d yolov5s.engine ../sample

如果要开启摄像头的话将yolov5.cpp的代码换成如下

  1. #include <iostream>
  2. #include <chrono>
  3. #include "cuda_utils.h"
  4. #include "logging.h"
  5. #include "common.hpp"
  6. #include "utils.h"
  7. #include "calibrator.h"
  8.  
  9. #define USE_FP16  // set USE_INT8 or USE_FP16 or USE_FP32
  10. #define DEVICE 0  // GPU id
  11. #define NMS_THRESH 0.4
  12. #define CONF_THRESH 0.5
  13. #define BATCH_SIZE 1
  14.  
  15. // stuff we know about the network and the input/output blobs
  16. static const int INPUT_H = Yolo::INPUT_H;
  17. static const int INPUT_W = Yolo::INPUT_W;
  18. static const int CLASS_NUM = Yolo::CLASS_NUM;
  19. static const int OUTPUT_SIZE = Yolo::MAX_OUTPUT_BBOX_COUNT * sizeof(Yolo::Detection) / sizeof(float) + 1;  // we assume the yololayer outputs no more than MAX_OUTPUT_BBOX_COUNT boxes that conf >= 0.1
  20. const char* INPUT_BLOB_NAME = "data";
  21. const char* OUTPUT_BLOB_NAME = "prob";
  22. static Logger gLogger;
  23.  
  24. //修改为自己的类别
  25. char *my_classes[]={ "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
  26.          "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
  27.          "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
  28.          "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard","surfboard",
  29.          "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
  30.          "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
  31.          "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
  32.          "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
  33.          "hair drier", "toothbrush" };
  34.  
  35. static int get_width(int x, float gw, int divisor = 8) {
  36.     //return math.ceil(x / divisor) * divisor
  37.     if (int(x * gw) % divisor == 0) {
  38.         return int(x * gw);
  39.     }
  40.     return (int(x * gw / divisor) + 1) * divisor;
  41. }
  42.  
  43. static int get_depth(int x, float gd) {
  44.     if (x == 1) {
  45.         return 1;
  46.     }
  47.     else {
  48.         return round(x * gd) > 1 ? round(x * gd) : 1;
  49.     }
  50. }
  51.  //#创建engine和network
  52. ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) {
  53.     INetworkDefinition* network = builder->createNetworkV2(0U);
  54.  
  55.     // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
  56.     ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W });
  57.     assert(data);
  58.  
  59.     std::map<std::string, Weights> weightMap = loadWeights(wts_name);
  60.  
  61.     /* ------ yolov5 backbone------ */
  62.     auto focus0 = focus(network, weightMap, *data, 3, get_width(64, gw), 3, "model.0");
  63.     auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1");
  64.     auto bottleneck_CSP2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2");
  65.     auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3");
  66.     auto bottleneck_csp4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(9, gd), true, 1, 0.5, "model.4");
  67.     auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5");
  68.     auto bottleneck_csp6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6");
  69.     auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.7");
  70.     auto spp8 = SPP(network, weightMap, *conv7->getOutput(0), get_width(1024, gw), get_width(1024, gw), 5, 9, 13, "model.8");
  71.  
  72.     /* ------ yolov5 head ------ */
  73.     auto bottleneck_csp9 = C3(network, weightMap, *spp8->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.9");
  74.     auto conv10 = convBlock(network, weightMap, *bottleneck_csp9->getOutput(0), get_width(512, gw), 1, 1, 1, "model.10");
  75.  
  76.     auto upsample11 = network->addResize(*conv10->getOutput(0));
  77.     assert(upsample11);
  78.     upsample11->setResizeMode(ResizeMode::kNEAREST);
  79.     upsample11->setOutputDimensions(bottleneck_csp6->getOutput(0)->getDimensions());
  80.  
  81.     ITensor* inputTensors12[] = { upsample11->getOutput(0), bottleneck_csp6->getOutput(0) };
  82.     auto cat12 = network->addConcatenation(inputTensors12, 2);
  83.     auto bottleneck_csp13 = C3(network, weightMap, *cat12->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.13");
  84.     auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), get_width(256, gw), 1, 1, 1, "model.14");
  85.  
  86.     auto upsample15 = network->addResize(*conv14->getOutput(0));
  87.     assert(upsample15);
  88.     upsample15->setResizeMode(ResizeMode::kNEAREST);
  89.     upsample15->setOutputDimensions(bottleneck_csp4->getOutput(0)->getDimensions());
  90.  
  91.     ITensor* inputTensors16[] = { upsample15->getOutput(0), bottleneck_csp4->getOutput(0) };
  92.     auto cat16 = network->addConcatenation(inputTensors16, 2);
  93.  
  94.     auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.17");
  95.  
  96.     // yolo layer 0
  97.     IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]);
  98.     auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256, gw), 3, 2, 1, "model.18");
  99.     ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) };
  100.     auto cat19 = network->addConcatenation(inputTensors19, 2);
  101.     auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.20");
  102.     //yolo layer 1
  103.     IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]);
  104.     auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), get_width(512, gw), 3, 2, 1, "model.21");
  105.     ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) };
  106.     auto cat22 = network->addConcatenation(inputTensors22, 2);
  107.     auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.23");
  108.     IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]);
  109.  
  110.     auto yolo = addYoLoLayer(network, weightMap, "model.24", std::vector<IConvolutionLayer*>{det0, det1, det2});
  111.     yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME);
  112.     network->markOutput(*yolo->getOutput(0));
  113.  
  114.     // Build engine
  115.     builder->setMaxBatchSize(maxBatchSize);
  116.     config->setMaxWorkspaceSize(16 * (1 << 20));  // 16MB
  117. #if defined(USE_FP16)
  118.     config->setFlag(BuilderFlag::kFP16);
  119. #elif defined(USE_INT8)
  120.     std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
  121.     assert(builder->platformHasFastInt8());
  122.     config->setFlag(BuilderFlag::kINT8);
  123.     Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME);
  124.     config->setInt8Calibrator(calibrator);
  125. #endif
  126.  
  127.     std::cout << "Building engine, please wait for a while..." << std::endl;
  128.     ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
  129.     std::cout << "Build engine successfully!" << std::endl;
  130.  
  131.     // Don't need the network any more
  132.     network->destroy();
  133.  
  134.     // Release host memory
  135.     for (auto& mem : weightMap)
  136.     {
  137.         free((void*)(mem.second.values));
  138.     }
  139.  
  140.     return engine;
  141. }
  142.  
  143. ICudaEngine* build_engine_p6(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) {
  144.     INetworkDefinition* network = builder->createNetworkV2(0U);
  145.  
  146.     // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
  147.     ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W });
  148.     assert(data);
  149.  
  150.     std::map<std::string, Weights> weightMap = loadWeights(wts_name);
  151.  
  152.     /* ------ yolov5 backbone------ */
  153.     auto focus0 = focus(network, weightMap, *data, 3, get_width(64, gw), 3, "model.0");
  154.     auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1");
  155.     auto c3_2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2");
  156.     auto conv3 = convBlock(network, weightMap, *c3_2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3");
  157.     auto c3_4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(9, gd), true, 1, 0.5, "model.4");
  158.     auto conv5 = convBlock(network, weightMap, *c3_4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5");
  159.     auto c3_6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6");
  160.     auto conv7 = convBlock(network, weightMap, *c3_6->getOutput(0), get_width(768, gw), 3, 2, 1, "model.7");
  161.     auto c3_8 = C3(network, weightMap, *conv7->getOutput(0), get_width(768, gw), get_width(768, gw), get_depth(3, gd), true, 1, 0.5, "model.8");
  162.     auto conv9 = convBlock(network, weightMap, *c3_8->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.9");
  163.     auto spp10 = SPP(network, weightMap, *conv9->getOutput(0), get_width(1024, gw), get_width(1024, gw), 3, 5, 7, "model.10");
  164.     auto c3_11 = C3(network, weightMap, *spp10->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.11");
  165.  
  166.     /* ------ yolov5 head ------ */
  167.     auto conv12 = convBlock(network, weightMap, *c3_11->getOutput(0), get_width(768, gw), 1, 1, 1, "model.12");
  168.     auto upsample13 = network->addResize(*conv12->getOutput(0));
  169.     assert(upsample13);
  170.     upsample13->setResizeMode(ResizeMode::kNEAREST);
  171.     upsample13->setOutputDimensions(c3_8->getOutput(0)->getDimensions());
  172.     ITensor* inputTensors14[] = { upsample13->getOutput(0), c3_8->getOutput(0) };
  173.     auto cat14 = network->addConcatenation(inputTensors14, 2);
  174.     auto c3_15 = C3(network, weightMap, *cat14->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.15");
  175.  
  176.     auto conv16 = convBlock(network, weightMap, *c3_15->getOutput(0), get_width(512, gw), 1, 1, 1, "model.16");
  177.     auto upsample17 = network->addResize(*conv16->getOutput(0));
  178.     assert(upsample17);
  179.     upsample17->setResizeMode(ResizeMode::kNEAREST);
  180.     upsample17->setOutputDimensions(c3_6->getOutput(0)->getDimensions());
  181.     ITensor* inputTensors18[] = { upsample17->getOutput(0), c3_6->getOutput(0) };
  182.     auto cat18 = network->addConcatenation(inputTensors18, 2);
  183.     auto c3_19 = C3(network, weightMap, *cat18->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.19");
  184.  
  185.     auto conv20 = convBlock(network, weightMap, *c3_19->getOutput(0), get_width(256, gw), 1, 1, 1, "model.20");
  186.     auto upsample21 = network->addResize(*conv20->getOutput(0));
  187.     assert(upsample21);
  188.     upsample21->setResizeMode(ResizeMode::kNEAREST);
  189.     upsample21->setOutputDimensions(c3_4->getOutput(0)->getDimensions());
  190.     ITensor* inputTensors21[] = { upsample21->getOutput(0), c3_4->getOutput(0) };
  191.     auto cat22 = network->addConcatenation(inputTensors21, 2);
  192.     auto c3_23 = C3(network, weightMap, *cat22->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.23");
  193.  
  194.     auto conv24 = convBlock(network, weightMap, *c3_23->getOutput(0), get_width(256, gw), 3, 2, 1, "model.24");
  195.     ITensor* inputTensors25[] = { conv24->getOutput(0), conv20->getOutput(0) };
  196.     auto cat25 = network->addConcatenation(inputTensors25, 2);
  197.     auto c3_26 = C3(network, weightMap, *cat25->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.26");
  198.  
  199.     auto conv27 = convBlock(network, weightMap, *c3_26->getOutput(0), get_width(512, gw), 3, 2, 1, "model.27");
  200.     ITensor* inputTensors28[] = { conv27->getOutput(0), conv16->getOutput(0) };
  201.     auto cat28 = network->addConcatenation(inputTensors28, 2);
  202.     auto c3_29 = C3(network, weightMap, *cat28->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.29");
  203.  
  204.     auto conv30 = convBlock(network, weightMap, *c3_29->getOutput(0), get_width(768, gw), 3, 2, 1, "model.30");
  205.     ITensor* inputTensors31[] = { conv30->getOutput(0), conv12->getOutput(0) };
  206.     auto cat31 = network->addConcatenation(inputTensors31, 2);
  207.     auto c3_32 = C3(network, weightMap, *cat31->getOutput(0), get_width(2048, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.32");
  208.  
  209.     /* ------ detect ------ */
  210.     IConvolutionLayer* det0 = network->addConvolutionNd(*c3_23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.0.weight"], weightMap["model.33.m.0.bias"]);
  211.     IConvolutionLayer* det1 = network->addConvolutionNd(*c3_26->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.1.weight"], weightMap["model.33.m.1.bias"]);
  212.     IConvolutionLayer* det2 = network->addConvolutionNd(*c3_29->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.2.weight"], weightMap["model.33.m.2.bias"]);
  213.     IConvolutionLayer* det3 = network->addConvolutionNd(*c3_32->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.3.weight"], weightMap["model.33.m.3.bias"]);
  214.  
  215.     auto yolo = addYoLoLayer(network, weightMap, "model.33", std::vector<IConvolutionLayer*>{det0, det1, det2, det3});
  216.     yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME);
  217.     network->markOutput(*yolo->getOutput(0));
  218.  
  219.     // Build engine
  220.     builder->setMaxBatchSize(maxBatchSize);
  221.     config->setMaxWorkspaceSize(16 * (1 << 20));  // 16MB
  222. #if defined(USE_FP16)
  223.     config->setFlag(BuilderFlag::kFP16);
  224. #elif defined(USE_INT8)
  225.     std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
  226.     assert(builder->platformHasFastInt8());
  227.     config->setFlag(BuilderFlag::kINT8);
  228.     Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME);
  229.     config->setInt8Calibrator(calibrator);
  230. #endif
  231.  
  232.     std::cout << "Building engine, please wait for a while..." << std::endl;
  233.     ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
  234.     std::cout << "Build engine successfully!" << std::endl;
  235.  
  236.     // Don't need the network any more
  237.     network->destroy();
  238.  
  239.     // Release host memory
  240.     for (auto& mem : weightMap)
  241.     {
  242.         free((void*)(mem.second.values));
  243.     }
  244.  
  245.     return engine;
  246. }
  247.  
  248. void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream, float& gd, float& gw, std::string& wts_name) {
  249.     // Create builder
  250.     IBuilder* builder = createInferBuilder(gLogger);
  251.     IBuilderConfig* config = builder->createBuilderConfig();
  252.  
  253.     // Create model to populate the network, then set the outputs and create an engine
  254.     ICudaEngine* engine = build_engine(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name);
  255.     assert(engine != nullptr);
  256.  
  257.     // Serialize the engine
  258.     (*modelStream) = engine->serialize();
  259.  
  260.     // Close everything down
  261.     engine->destroy();
  262.     builder->destroy();
  263.     config->destroy();
  264. }
  265.  
  266. void doInference(IExecutionContext& context, cudaStream_t& stream, void** buffers, float* input, float* output, int batchSize) {
  267.     // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
  268.     CUDA_CHECK(cudaMemcpyAsync(buffers[0], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
  269.     context.enqueue(batchSize, buffers, stream, nullptr);
  270.     CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
  271.     cudaStreamSynchronize(stream);
  272. }
  273.  
  274. bool parse_args(int argc, char** argv, std::string& engine) {
  275.     if (argc < 3) return false;
  276.     if (std::string(argv[1]) == "-v" && argc == 3) {
  277.         engine = std::string(argv[2]);
  278.     }
  279.     else {
  280.         return false;
  281.     }
  282.     return true;
  283. }
  284.  
  285. int main(int argc, char** argv) {
  286.     cudaSetDevice(DEVICE);
  287.  
  288.     //std::string wts_name = "";
  289.     std::string engine_name = "";
  290.     //float gd = 0.0f, gw = 0.0f;
  291.     //std::string img_dir;
  292.  
  293.     if (!parse_args(argc, argv, engine_name)) {
  294.         std::cerr << "arguments not right!" << std::endl;
  295.         std::cerr << "./yolov5 -v [.engine] // run inference with camera" << std::endl;
  296.         return -1;
  297.     }
  298.  
  299.     std::ifstream file(engine_name, std::ios::binary);
  300.     if (!file.good()) {
  301.         std::cerr << " read " << engine_name << " error! " << std::endl;
  302.         return -1;
  303.     }
  304.     char* trtModelStream{ nullptr };
  305.     size_t size = 0;
  306.     file.seekg(0, file.end);
  307.     size = file.tellg();
  308.     file.seekg(0, file.beg);
  309.     trtModelStream = new char[size];
  310.     assert(trtModelStream);
  311.     file.read(trtModelStream, size);
  312.     file.close();
  313.  
  314.  
  315.     // prepare input data ---------------------------
  316.     static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W];
  317.     //for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
  318.     //    data[i] = 1.0;
  319.     static float prob[BATCH_SIZE * OUTPUT_SIZE];
  320.     IRuntime* runtime = createInferRuntime(gLogger);
  321.     assert(runtime != nullptr);
  322.     ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
  323.     assert(engine != nullptr);
  324.     IExecutionContext* context = engine->createExecutionContext();
  325.     assert(context != nullptr);
  326.     delete[] trtModelStream;
  327.     assert(engine->getNbBindings() == 2);
  328.     void* buffers[2];
  329.     // In order to bind the buffers, we need to know the names of the input and output tensors.
  330.     // Note that indices are guaranteed to be less than IEngine::getNbBindings()
  331.     const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME);
  332.     const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME);
  333.     assert(inputIndex == 0);
  334.     assert(outputIndex == 1);
  335.     // Create GPU buffers on device
  336.     CUDA_CHECK(cudaMalloc(&buffers[inputIndex], BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float)));
  337.     CUDA_CHECK(cudaMalloc(&buffers[outputIndex], BATCH_SIZE * OUTPUT_SIZE * sizeof(float)));
  338.     // Create stream
  339.     cudaStream_t stream;
  340.     CUDA_CHECK(cudaStreamCreate(&stream));
  341.  
  342.      //#读取本地视频
  343.     //cv::VideoCapture capture("/home/nano/Videos/video.mp4");
  344.      //#调用本地usb摄像头,我的默认参数为1,如果1报错,可修改为0.
  345.     cv::VideoCapture capture(1);
  346.     if (!capture.isOpened()) {
  347.         std::cout << "Error opening video stream or file" << std::endl;
  348.         return -1;
  349.     }
  350.  
  351.     int key;
  352.     int fcount = 0;
  353.     while (1)
  354.     {
  355.         cv::Mat frame;
  356.         capture >> frame;
  357.         if (frame.empty())
  358.         {
  359.             std::cout << "Fail to read image from camera!" << std::endl;
  360.             break;
  361.         }
  362.         fcount++;
  363.         //if (fcount < BATCH_SIZE && f + 1 != (int)file_names.size()) continue;
  364.         for (int b = 0; b < fcount; b++) {
  365.             //cv::Mat img = cv::imread(img_dir + "/" + file_names[f - fcount + 1 + b]);
  366.             cv::Mat img = frame;
  367.             if (img.empty()) continue;
  368.             cv::Mat pr_img = preprocess_img(img, INPUT_W, INPUT_H); // letterbox BGR to RGB
  369.             int i = 0;
  370.             for (int row = 0; row < INPUT_H; ++row) {
  371.                 uchar* uc_pixel = pr_img.data + row * pr_img.step;
  372.                 for (int col = 0; col < INPUT_W; ++col) {
  373.                     data[b * 3 * INPUT_H * INPUT_W + i] = (float)uc_pixel[2] / 255.0;
  374.                     data[b * 3 * INPUT_H * INPUT_W + i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0;
  375.                     data[b * 3 * INPUT_H * INPUT_W + i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.0;
  376.                     uc_pixel += 3;
  377.                     ++i;
  378.                 }
  379.             }
  380.         }
  381.  
  382.         // Run inference
  383.         auto start = std::chrono::system_clock::now();//#获取模型推理开始时间
  384.         doInference(*context, stream, buffers, data, prob, BATCH_SIZE);
  385.         auto end = std::chrono::system_clock::now();//#结束时间
  386.         //std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
  387.         int fps = 1000.0 / std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
  388.         std::vector<std::vector<Yolo::Detection>> batch_res(fcount);
  389.         for (int b = 0; b < fcount; b++) {
  390.             auto& res = batch_res[b];
  391.             nms(res, &prob[b * OUTPUT_SIZE], CONF_THRESH, NMS_THRESH);
  392.         }
  393.         for (int b = 0; b < fcount; b++) {
  394.             auto& res = batch_res[b];
  395.             //std::cout << res.size() << std::endl;
  396.             //cv::Mat img = cv::imread(img_dir + "/" + file_names[f - fcount + 1 + b]);
  397.             for (size_t j = 0; j < res.size(); j++) {
  398.                 cv::Rect r = get_rect(frame, res[j].bbox);
  399.                 cv::rectangle(frame, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
  400.                 std::string label = my_classes[(int)res[j].class_id];
  401.                 cv::putText(frame, label, cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
  402.                 std::string jetson_fps = "FPS: " + std::to_string(fps);
  403.                 cv::putText(frame, jetson_fps, cv::Point(11, 80), cv::FONT_HERSHEY_PLAIN, 3, cv::Scalar(0, 0, 255), 2, cv::LINE_AA);
  404.             }
  405.             //cv::imwrite("_" + file_names[f - fcount + 1 + b], img);
  406.         }
  407.         cv::imshow("yolov5", frame);
  408.         key = cv::waitKey(1);
  409.         if (key == 'q') {
  410.             break;
  411.         }
  412.         fcount = 0;
  413.     }
  414.  
  415.     capture.release();
  416.     // Release stream and buffers
  417.     cudaStreamDestroy(stream);
  418.     CUDA_CHECK(cudaFree(buffers[inputIndex]));
  419.     CUDA_CHECK(cudaFree(buffers[outputIndex]));
  420.     // Destroy the engine
  421.     context->destroy();
  422.     engine->destroy();
  423.     runtime->destroy();
  424.  
  425.     return 0;
  426. }

修改完代码后

  1. make
  2. sudo ./yolov5 -v yolov5s.engine
声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/小小林熬夜学编程/article/detail/449689
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号