yolov5s模型转tensorrt+deepstream检测+CSI和USB摄像头检测

软硬件版本

ubuntu 18.04

Jetson Nano 4G
JP 4.5.1
CUDA 10.2
TensorRT 7.1.3.0

python 3.6

前言

把自己训练好的yolov5s.pt转为tensorrt模型是为了实现推理加速，过程也非常简单，不过有一点就是网上的教程很多但是yolov5的代码版本也有很多，所以导致很多操作有一些不同，所以建议直接去大佬的开源看他的readme，下面我就来说一下我yolov5-5.0的转换过程

以下步骤都是在jetson nano上进行

一、安装pycuda

sudo pip3 install pycuda

这里我直接安装出现了错误

安装不上，于是我们只能进行自己编译安装

安装前先配置环境

$ sudo nano ~/.bashrc
export PATH=/usr/local/cuda-10.2/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
export CUDA_HOME=$CUDA_HOME:/usr/local/cuda-10.2
$ sudo source ~/.bashrc
$ nvcc -V 检测一下是否配置成功

我们首先下载https://files.pythonhosted.org/packages/5e/3f/5658c38579b41866ba21ee1b5020b8225cec86fe717e4b1c5c972de0a33c/pycuda-2019.1.2.tar.gz

然后进行解压安装一系列操作

我这里是在电脑上解压的然后把解压出来的文件夹传到了nano上

cd pycuda-2019.1.2/   #这一步是进入解压出来的文件夹里面
python3 configure.py --cuda-root=/usr/local/cuda-10.2  #下面这两步是进行编译安装
sudo python3 setup.py install

安装成功显示

详细步骤大家也可以参考：

jetson nano安装pycuda！！！_帅的发光发亮的博客-CSDN博客

二、克隆tensorrtx源码和yolov5源码

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

git clone -b v5.0 https://github.com/ultralytics/yolov5.git

三、生成yolov5s.wts文件

1.进入tensorrtx文件夹下的yolov5文件再，找到gen_wts.py，并把它复制到我们刚才下载的yolov5-5.0中

2.

cd yolov5-5.0 #进入这个文件夹下

3.把我们训练好的best.pt尽量改名为yolov5s.pt并复制到yolov5-5.0根目录下

输入如下命令生成wts文件

sudo python3 gen_wts.py -w yolov5s.pt -o yolov5s.wts

现在就在yolov5-5.0中生成一个yolov5s.wts文件

四、生成engine文件

1.进入tenorrtx/yolov5文件夹中进行编译

先将yololayer.h里的CLASS_NUM修改成你的。因为官方用的是coco数据集，所以默认是80。

然后开始编译

mkdir build     #创建build文件夹
将第三步生成的.wts文件复制到tensorrtx/yolov5/build里。
cd build
cmake ..       #一定要注意是cmake空格..
make           #每次修改为CLASS_NUM都要make一次

 make成功

2. 将第三步生成的.wts文件复制到tensorrtx/yolov5里。

3.生成.engine文件（我用的是yolov5s，所以结尾用s）

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

4.用它自带的图片测试一下

sudo ./yolov5 -d yolov5s.engine ../samples

接下来进行yolov5s.engine的调用USB摄像头检测（不使用deepstream）！！！

1.我们上面已经生成了yolov5s.engine，现在对yolov5.cpp进行魔改（注意是先要原来的yolo.cpp生成engine文件，然后再修改编译调用USB）

以下代码的借用这位大哥改好的，Jetson nano + yolov5 + TensorRT加速+调用usb摄像头_hahasl555的博客-CSDN博客

打开yolov5.cpp

2.把下面这个代码直接替换原来的yolov5.cpp

#include <iostream>
#include <chrono>
#include "cuda_utils.h"
#include "logging.h"
#include "common.hpp"
#include "utils.h"
#include "calibrator.h"

#define USE_FP16  // set USE_INT8 or USE_FP16 or USE_FP32
#define DEVICE 0  // GPU id
#define NMS_THRESH 0.4
#define CONF_THRESH 0.5
#define BATCH_SIZE 1

// stuff we know about the network and the input/output blobs
static const int INPUT_H = Yolo::INPUT_H;
static const int INPUT_W = Yolo::INPUT_W;
static const int CLASS_NUM = Yolo::CLASS_NUM;
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
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
static Logger gLogger;

char* my_classes[] = { "person" , "dog" , "bicycle" };

static int get_width(int x, float gw, int divisor = 8) {
    //return math.ceil(x / divisor) * divisor
    if (int(x * gw) % divisor == 0) {
        return int(x * gw);
    }
    return (int(x * gw / divisor) + 1) * divisor;
}

static int get_depth(int x, float gd) {
    if (x == 1) {
        return 1;
    }
    else {
        return round(x * gd) > 1 ? round(x * gd) : 1;
    }
}

ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) {
    INetworkDefinition* network = builder->createNetworkV2(0U);

    // Create input tensor of shape {3, INPUT_H, INPUT_W} with name INPUT_BLOB_NAME
    ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ 3, INPUT_H, INPUT_W });
    assert(data);

    std::map<std::string, Weights> weightMap = loadWeights(wts_name);

    /* ------ yolov5 backbone------ */
    auto focus0 = focus(network, weightMap, *data, 3, get_width(64, gw), 3, "model.0");
    auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1");
    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