ASCII码 ASCII码

Jetson Nano配置YOLOv5并实现FPS=25

发布于：2022-05-08 01:06:35  栏目：技术文档

镜像下载、域名解析、时间同步请点击 阿里云开源镜像站

一、版本说明

JetPack 4.6——2021.8yolov5-v6.0版本使用的为yolov5的yolov5n.pt，并利用tensorrtx进行加速推理，在调用摄像头实时检测可以达到FPS=25。

二、配置CUDA

  1. sudo gedit ~/.bashrc

在打开的文档的末尾添加如下：

  1. export CUDA_HOME=/usr/local/cuda-10.2
  2. export LD_LIBRARY_PATH=/usr/local/cuda-10.2/lib64:$LD_LIBRARY_PATH
  3. export PATH=/usr/local/cuda-10.2/bin:$PATH

保持并退出，终端执行

  1. source ~/.bashrc
  2. nvcc -V    #如果配置成功可以看到CUDA的版本号

三、修改Nano板显存

1.打开终端输入：

  1. sudo gedit /etc/systemd/nvzramconfig.sh

2.修改nvzramconfig.sh文件

  1. 修改mem = $((("${totalmem}"/2/"${NRDEVICES}")*1024))
  2. 为mem = $((("${totalmem}"*2/"${NRDEVICES}")*1024))

3.重启Jetson Nano

4.终端中输入：

  1. free -h

可查看到swap已经变为7.7G

四、配置Pytorch1.8

1.下载torch-1.8.0-cp36-cp36m-linux_aarch64.whl

下载地址：nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl

说明：建议在电脑上下载后拷贝到Jetson Nano的文件夹下，因为该网站的服务器在国外，可能下载比较慢或网页加载不出来，可以打开VPN进行下载。

2.安装所需的依赖包及pytorch

打开终端输入：

  1. sudo apt-get update
  2. sudo apt-get upgrade
  3. sudo apt-get dist-upgrade
  4. sudo apt-get install python3-pip libopenblas-base libopenmpi-dev

因为下面用pip指令安装时用默认选用的国外源，所以下载比较费时间，建议更换一下国内源，这里提供阿里源，当使用某种国内源pip无法下载某一包时可以尝试切换再下载。具体步骤如下：

打开终端输入：

  1. mkdir ~/.pip
  2. sudo gedit ~/.pip/pip.conf

在空白文件中输入如下内容保存并退出：

以下为阿里源

  1. [global]
  2. index-url=http://mirrors.aliyun.com/pypi/simple/
  3. [install]
  4. trusted-host=mirrors.aliyun.com

终端输入：

  1. pip3 install --upgrade pip        #如果pip已是最新，可不执行
  2. pip3 install Cython
  3. pip3 install numpy
  4. pip3 install torch-1.8.0-cp36-cp36m-linux_aarch64.whl        #注意要在存放该文件下的位置打开终端并运行
  5. sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev
  6. git clone --branch v0.9.0 https://github.com/pytorch/vision torchvision  #下载torchvision，会下载一个文件夹
  7. cd torchvision    #或者进入到这个文件夹，右键打开终端
  8. export BUILD_VERSION=0.9.0
  9. python3 setup.py install --user    #时间较久
  10. #验证torch和torchvision这两个模块是否安装成功
  11. python3
  12. import torch
  13. print(torch.__version__)    #注意version前后都是有两个横杠
  14. #如果安装成功会打印出版本号
  15. import torchvision
  16. print(torchvision.__version__)
  17. #如果安装成功会打印出版本号

五、搭建yolov5环境

终端中输入：

  1. git clone https://github.com/ultralytics/yolov5.git        #因为不开VPN很容易下载出错，建议在电脑中下载后拷贝到jetson nano中
  2. python3 -m pip install --upgrade pip
  3. cd yolov5    #如果是手动下载的，文件名称为yolov5-master.zip压缩包格式，所以要对用unzip yolov5-master.zip进行解压，然后再进入到该文件夹
  4. pip3 install -r requirements.txt        #我的问题是对matplotlib包装不上，解决办法，在下方。如果其他包安装不上，可去重新执行换源那一步，更换另一种国内源。
  5. python3 -m pip list        #可查看python中安装的包
  6. 以下指令可以用来测试yolov5
  7. python3 detect.py --source data/images/bus.jpg --weights yolov5n.pt --img 640    #图片测试
  8. python3 detect.py --source video.mp4 --weights yolov5n.pt --img 640    #视频测试,需要自己准备视频
  9. python3 detect.py --source 0 --weights yolov5n.pt --img 640    #摄像头测试

问题1：解决matplotlib安装不上问题

解决：下载matplotlib的whl包（下方有网盘分享）

问题2：在运行yolov5的detect.py文件时出现 “Illegal instruction（core dumped）”

解决：

  1. sudo gedit ~/.bashrc
  2. 末尾添加
  3. export OPENBLAS_CORETYPE=ARMV8
  4. 保持关闭
  5. source ~/.bashrc

file

file

六、利用tensorrtx加速推理

1.下载tensorrtx

下载地址：https://github.com/wang-xinyu/tensorrtx.git

或者

  1. git clone https://github.com/wang-xinyu/tensorrtx.git

2.编译

将下载的tensorrtx项目中的yolov5/gen_wts.py复制到上述的yolov5（注意：不是tensorrtx下的yolov5！！！）下，然后在此处打开终端

打开终端输入：

  1. python3 gen_wts.py -w yolov5n.pt -o yolov5n.wts        #生成wts文件，要先把yolov5n.pt文件放在此处再去执行
  2. cd ~/tensorrtx/yolov5/        #如果是手动下载的名称可能是tensorrtx-master
  3. mkdir build
  4. cd build
  5. 将生成的wts文件复制到build下    #手动下载的，名称为yolov5-master
  6. cmake ..
  7. make -j4
  8. sudo ./yolov5 -s yolov5n.wts yolov5n.engine n #生成engine文件
  9. sudo ./yolov5 -d yolov5n.engine ../samples/    #测试图片查看效果,发现在检测zidane.jpg时漏检，这时可以返回上一层文件夹找到yolov5.cpp中的CONF_THRESH=0.25再进入到build中重新make -j4，再重新运行该指令即可

3.调用USB摄像头

参考了该文章https://blog.csdn.net/weixin_54603153/article/details/120079220

（1）在tensorrtx/yolov5下备份yolov5.cpp文件，因为如果更换模型时重新推理加速时需要用到该文件。

（2）然后对yolov5.cpp文件修改为如下内容

修改了12行和342行

  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"
  9. #define USE_FP32  // set USE_INT8 or USE_FP16 or USE_FP32
  10. #define DEVICE 0  // GPU id
  11. #define NMS_THRESH 0.4 //0.4
  12. #define CONF_THRESH 0.25    //置信度，默认值为0.5，由于效果不好修改为0.25取得了较好的效果
  13. #define BATCH_SIZE 1
  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;
  24. char* my_classes[] = { "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
  25.          "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
  26.          "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
  27.          "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard","surfboard",
  28.          "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
  29.          "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
  30.          "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
  31.          "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
  32.          "hair drier", "toothbrush" };
  34. static int get_width(int x, float gw, int divisor = 8) {
  35.     //return math.ceil(x / divisor) * divisor
  36.     if (int(x * gw) % divisor == 0) {
  37.         return int(x * gw);
  38.     }
  39.     return (int(x * gw / divisor) + 1) * divisor;
  40. }
  42. static int get_depth(int x, float gd) {
  43.     if (x == 1) {
  44.         return 1;
  45.     }
  46.     else {
  47.         return round(x * gd) > 1 ? round(x * gd) : 1;
  48.     }
  49. }
  51. ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) {
  52.     INetworkDefinition* network = builder->createNetworkV2(0U);
  54.     // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
  55.     ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W });
  56.     assert(data);
  58.     std::map<std::string, Weights> weightMap = loadWeights(wts_name);
  60.     /* ------ yolov5 backbone------ */
  61.     auto focus0 = focus(network, weightMap, *data, 3, get_width(64, gw), 3, "model.0");
  62.     auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1");
  63.     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");
  64.     auto conv3 = convBlock(network, weightMap, *bottleneck_CSP2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3");
  65.     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");
  66.     auto conv5 = convBlock(network, weightMap, *bottleneck_csp4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5");
  67.     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");
  68.     auto conv7 = convBlock(network, weightMap, *bottleneck_csp6->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.7");
  69.     auto spp8 = SPP(network, weightMap, *conv7->getOutput(0), get_width(1024, gw), get_width(1024, gw), 5, 9, 13, "model.8");
  71.     /* ------ yolov5 head ------ */
  72.     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");
  73.     auto conv10 = convBlock(network, weightMap, *bottleneck_csp9->getOutput(0), get_width(512, gw), 1, 1, 1, "model.10");
  75.     auto upsample11 = network->addResize(*conv10->getOutput(0));
  76.     assert(upsample11);
  77.     upsample11->setResizeMode(ResizeMode::kNEAREST);
  78.     upsample11->setOutputDimensions(bottleneck_csp6->getOutput(0)->getDimensions());
  80.     ITensor* inputTensors12[] = { upsample11->getOutput(0), bottleneck_csp6->getOutput(0) };
  81.     auto cat12 = network->addConcatenation(inputTensors12, 2);
  82.     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");
  83.     auto conv14 = convBlock(network, weightMap, *bottleneck_csp13->getOutput(0), get_width(256, gw), 1, 1, 1, "model.14");
  85.     auto upsample15 = network->addResize(*conv14->getOutput(0));
  86.     assert(upsample15);
  87.     upsample15->setResizeMode(ResizeMode::kNEAREST);
  88.     upsample15->setOutputDimensions(bottleneck_csp4->getOutput(0)->getDimensions());
  90.     ITensor* inputTensors16[] = { upsample15->getOutput(0), bottleneck_csp4->getOutput(0) };
  91.     auto cat16 = network->addConcatenation(inputTensors16, 2);
  93.     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.     // yolo layer 0
  96.     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"]);
  97.     auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256, gw), 3, 2, 1, "model.18");
  98.     ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) };
  99.     auto cat19 = network->addConcatenation(inputTensors19, 2);
  100.     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");
  101.     //yolo layer 1
  102.     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"]);
  103.     auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), get_width(512, gw), 3, 2, 1, "model.21");
  104.     ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) };
  105.     auto cat22 = network->addConcatenation(inputTensors22, 2);
  106.     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");
  107.     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.     auto yolo = addYoLoLayer(network, weightMap, "model.24", std::vector<IConvolutionLayer*>{det0, det1, det2});
  110.     yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME);
  111.     network->markOutput(*yolo->getOutput(0));
  113.     // Build engine
  114.     builder->setMaxBatchSize(maxBatchSize);
  115.     config->setMaxWorkspaceSize(16 * (1 << 20));  // 16MB
  116. #if defined(USE_FP16)
  117.     config->setFlag(BuilderFlag::kFP16);
  118. #elif defined(USE_INT8)
  119.     std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
  120.     assert(builder->platformHasFastInt8());
  121.     config->setFlag(BuilderFlag::kINT8);
  122.     Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME);
  123.     config->setInt8Calibrator(calibrator);
  124. #endif
  126.     std::cout << "Building engine, please wait for a while..." << std::endl;
  127.     ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
  128.     std::cout << "Build engine successfully!" << std::endl;
  130.     // Don't need the network any more
  131.     network->destroy();
  133.     // Release host memory
  134.     for (auto& mem : weightMap)
  135.     {
  136.         free((void*)(mem.second.values));
  137.     }
  139.     return engine;
  140. }
  142. ICudaEngine* build_engine_p6(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) {
  143.     INetworkDefinition* network = builder->createNetworkV2(0U);
  145.     // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
  146.     ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W });
  147.     assert(data);
  149.     std::map<std::string, Weights> weightMap = loadWeights(wts_name);
  151.     /* ------ yolov5 backbone------ */
  152.     auto focus0 = focus(network, weightMap, *data, 3, get_width(64, gw), 3, "model.0");
  153.     auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1");
  154.     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");
  155.     auto conv3 = convBlock(network, weightMap, *c3_2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3");
  156.     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");
  157.     auto conv5 = convBlock(network, weightMap, *c3_4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5");
  158.     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");
  159.     auto conv7 = convBlock(network, weightMap, *c3_6->getOutput(0), get_width(768, gw), 3, 2, 1, "model.7");
  160.     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");
  161.     auto conv9 = convBlock(network, weightMap, *c3_8->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.9");
  162.     auto spp10 = SPP(network, weightMap, *conv9->getOutput(0), get_width(1024, gw), get_width(1024, gw), 3, 5, 7, "model.10");
  163.     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.     /* ------ yolov5 head ------ */
  166.     auto conv12 = convBlock(network, weightMap, *c3_11->getOutput(0), get_width(768, gw), 1, 1, 1, "model.12");
  167.     auto upsample13 = network->addResize(*conv12->getOutput(0));
  168.     assert(upsample13);
  169.     upsample13->setResizeMode(ResizeMode::kNEAREST);
  170.     upsample13->setOutputDimensions(c3_8->getOutput(0)->getDimensions());
  171.     ITensor* inputTensors14[] = { upsample13->getOutput(0), c3_8->getOutput(0) };
  172.     auto cat14 = network->addConcatenation(inputTensors14, 2);
  173.     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.     auto conv16 = convBlock(network, weightMap, *c3_15->getOutput(0), get_width(512, gw), 1, 1, 1, "model.16");
  176.     auto upsample17 = network->addResize(*conv16->getOutput(0));
  177.     assert(upsample17);
  178.     upsample17->setResizeMode(ResizeMode::kNEAREST);
  179.     upsample17->setOutputDimensions(c3_6->getOutput(0)->getDimensions());
  180.     ITensor* inputTensors18[] = { upsample17->getOutput(0), c3_6->getOutput(0) };
  181.     auto cat18 = network->addConcatenation(inputTensors18, 2);
  182.     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.     auto conv20 = convBlock(network, weightMap, *c3_19->getOutput(0), get_width(256, gw), 1, 1, 1, "model.20");
  185.     auto upsample21 = network->addResize(*conv20->getOutput(0));
  186.     assert(upsample21);
  187.     upsample21->setResizeMode(ResizeMode::kNEAREST);
  188.     upsample21->setOutputDimensions(c3_4->getOutput(0)->getDimensions());
  189.     ITensor* inputTensors21[] = { upsample21->getOutput(0), c3_4->getOutput(0) };
  190.     auto cat22 = network->addConcatenation(inputTensors21, 2);
  191.     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.     auto conv24 = convBlock(network, weightMap, *c3_23->getOutput(0), get_width(256, gw), 3, 2, 1, "model.24");
  194.     ITensor* inputTensors25[] = { conv24->getOutput(0), conv20->getOutput(0) };
  195.     auto cat25 = network->addConcatenation(inputTensors25, 2);
  196.     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.     auto conv27 = convBlock(network, weightMap, *c3_26->getOutput(0), get_width(512, gw), 3, 2, 1, "model.27");
  199.     ITensor* inputTensors28[] = { conv27->getOutput(0), conv16->getOutput(0) };
  200.     auto cat28 = network->addConcatenation(inputTensors28, 2);
  201.     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.     auto conv30 = convBlock(network, weightMap, *c3_29->getOutput(0), get_width(768, gw), 3, 2, 1, "model.30");
  204.     ITensor* inputTensors31[] = { conv30->getOutput(0), conv12->getOutput(0) };
  205.     auto cat31 = network->addConcatenation(inputTensors31, 2);
  206.     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.     /* ------ detect ------ */
  209.     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"]);
  210.     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"]);
  211.     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"]);
  212.     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.     auto yolo = addYoLoLayer(network, weightMap, "model.33", std::vector<IConvolutionLayer*>{det0, det1, det2, det3});
  215.     yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME);
  216.     network->markOutput(*yolo->getOutput(0));
  218.     // Build engine
  219.     builder->setMaxBatchSize(maxBatchSize);
  220.     config->setMaxWorkspaceSize(16 * (1 << 20));  // 16MB
  221. #if defined(USE_FP16)
  222.     config->setFlag(BuilderFlag::kFP16);
  223. #elif defined(USE_INT8)
  224.     std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
  225.     assert(builder->platformHasFastInt8());
  226.     config->setFlag(BuilderFlag::kINT8);
  227.     Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME);
  228.     config->setInt8Calibrator(calibrator);
  229. #endif
  231.     std::cout << "Building engine, please wait for a while..." << std::endl;
  232.     ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
  233.     std::cout << "Build engine successfully!" << std::endl;
  235.     // Don't need the network any more
  236.     network->destroy();
  238.     // Release host memory
  239.     for (auto& mem : weightMap)
  240.     {
  241.         free((void*)(mem.second.values));
  242.     }
  244.     return engine;
  245. }
  247. void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream, float& gd, float& gw, std::string& wts_name) {
  248.     // Create builder
  249.     IBuilder* builder = createInferBuilder(gLogger);
  250.     IBuilderConfig* config = builder->createBuilderConfig();
  252.     // Create model to populate the network, then set the outputs and create an engine
  253.     ICudaEngine* engine = build_engine(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name);
  254.     assert(engine != nullptr);
  256.     // Serialize the engine
  257.     (*modelStream) = engine->serialize();
  259.     // Close everything down
  260.     engine->destroy();
  261.     builder->destroy();
  262.     config->destroy();
  263. }
  265. void doInference(IExecutionContext& context, cudaStream_t& stream, void** buffers, float* input, float* output, int batchSize) {
  266.     // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
  267.     CUDA_CHECK(cudaMemcpyAsync(buffers[0], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
  268.     context.enqueue(batchSize, buffers, stream, nullptr);
  269.     CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
  270.     cudaStreamSynchronize(stream);
  271. }
  273. bool parse_args(int argc, char** argv, std::string& engine) {
  274.     if (argc < 3) return false;
  275.     if (std::string(argv[1]) == "-v" && argc == 3) {
  276.         engine = std::string(argv[2]);
  277.     }
  278.     else {
  279.         return false;
  280.     }
  281.     return true;
  282. }
  284. int main(int argc, char** argv) {
  285.     cudaSetDevice(DEVICE);
  287.     //std::string wts_name = "";
  288.     std::string engine_name = "";
  289.     //float gd = 0.0f, gw = 0.0f;
  290.     //std::string img_dir;
  292.     if (!parse_args(argc, argv, engine_name)) {
  293.         std::cerr << "arguments not right!" << std::endl;
  294.         std::cerr << "./yolov5 -v [.engine] // run inference with camera" << std::endl;
  295.         return -1;
  296.     }
  298.     std::ifstream file(engine_name, std::ios::binary);
  299.     if (!file.good()) {
  300.         std::cerr << " read " << engine_name << " error! " << std::endl;
  301.         return -1;
  302.     }
  303.     char* trtModelStream{ nullptr };
  304.     size_t size = 0;
  305.     file.seekg(0, file.end);
  306.     size = file.tellg();
  307.     file.seekg(0, file.beg);
  308.     trtModelStream = new char[size];
  309.     assert(trtModelStream);
  310.     file.read(trtModelStream, size);
  311.     file.close();
  314.     // prepare input data ---------------------------
  315.     static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W];
  316.     //for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
  317.     //    data[i] = 1.0;
  318.     static float prob[BATCH_SIZE * OUTPUT_SIZE];
  319.     IRuntime* runtime = createInferRuntime(gLogger);
  320.     assert(runtime != nullptr);
  321.     ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
  322.     assert(engine != nullptr);
  323.     IExecutionContext* context = engine->createExecutionContext();
  324.     assert(context != nullptr);
  325.     delete[] trtModelStream;
  326.     assert(engine->getNbBindings() == 2);
  327.     void* buffers[2];
  328.     // In order to bind the buffers, we need to know the names of the input and output tensors.
  329.     // Note that indices are guaranteed to be less than IEngine::getNbBindings()
  330.     const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME);
  331.     const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME);
  332.     assert(inputIndex == 0);
  333.     assert(outputIndex == 1);
  334.     // Create GPU buffers on device
  335.     CUDA_CHECK(cudaMalloc(&buffers[inputIndex], BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float)));
  336.     CUDA_CHECK(cudaMalloc(&buffers[outputIndex], BATCH_SIZE * OUTPUT_SIZE * sizeof(float)));
  337.     // Create stream
  338.     cudaStream_t stream;
  339.     CUDA_CHECK(cudaStreamCreate(&stream));
  342.     cv::VideoCapture capture("/home/cao-yolox/yolov5/tensorrtx-master/yolov5/samples/1.mp4");    #修改为自己要检测的视频或者图片，注意要写全路径，如果调用摄像头，则括号内的参数设为0，注意引号要去掉。
  343.     //cv::VideoCapture capture("../overpass.mp4");
  344.     //int fourcc = cv::VideoWriter::fourcc('M','J','P','G');
  345.     //capture.set(cv::CAP_PROP_FOURCC, fourcc);
  346.     if (!capture.isOpened()) {
  347.         std::cout << "Error opening video stream or file" << std::endl;
  348.         return -1;
  349.     }
  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.         }
  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 = "Jetson Nano 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.     }
  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();
  425.     return 0;
  426. }

4.重新编译

进入到buid下重新make。注意只要修改了yolov5.cpp就要重新make。执行

  1. sudo ./yolov5 -v yolov5n.engine    #注意要提前插好摄像头

问题：出现Failed to load module “canberra-gtk-module”

解决：

  1. sudo apt-get install libcanberra-gtk-module

5.效果

如下的测试，是在一个公用的行人检测的视频上进行的，如果想用可在如下链接下载：

链接：https://pan.baidu.com/s/1HivF1OifVA8pHnGKtkXPfg

提取码：jr7o

file

本文转自：https://blog.csdn.net/carrymingteng/article/details/120978053

相关推荐
阅读 +

© 2020 asciim码

人生就是一场修行