Jetson nano + yolov5 + TensorRT加速+调用usb摄像头
目录
- 前言
- 一、环境安装
-
- 1、安装虚拟环境virtualenv（可选）
- 2、设置cuda环境变量，解决nvcc -V找不到命令
- 3、更新
- 4、安装pytorch 和 torchvision
- 5、安装yolov5必须环境
- 二、TensorRt加速
- 三、调用usb摄像头
- 总结
前言
因为在做毕业设计，需要将yolov5移植到嵌入式开发板。以前在Firefly-RK3399的Tengine框架上部署过Mobilenet-SSD，也玩过树莓派。这次尝试使用Jetson nano部署yolov5（4.0版本）。可惜更新换代太快，网上的教程比较旧，上个月又出了yolov5(5.0版)，踩了很多坑，让我下定决心要更新一版教程。
历时半个月，尝试了网上大多数方案，重装了八次，总结出以下最优方案。（2021-5-8）
本文忽略镜像烧写过程，因为太简单了。
一、环境安装
1、安装虚拟环境virtualenv（可选）
我除了第一次害怕环境出错外，后面几次都没有安装。主要是怕麻烦。
sudo pip3 install virtualenv virtualenvwrapper
修改环境变量
nano ~/.bashrc
添加下面三行到最下面
export WORKON_HOME=$HOME/.virtualenvs
export VIRTUALENVWRAPPER_PYTHON=/usr/bin/python3
export /usr/local/bin/virtualenvwrapper.sh
接着创建并进入虚拟环境，与原生的OpenCV建立链接
#创建一个叫yolov5的环境，你也可以叫其他名：
mkvirtualenv yolov5
#激活环境：
workon yolov5
#找到原生的OpenCV位置（寻找 .so 档案，大家应该都一样会在 /usr/lib/python3.6/dist-packages/cv2/python-3.6/ 里面）：
sudo find / -name cv2
#建立链接
ln -s /usr/lib/python3.6/dist-packages/cv2/python-3.6/cv2.cpython-36m-aarch64-linux-gnu.so ./virtualenvs/你的虚拟环境名/lib/python3.6/site-packages
附上virtualenv常用命令
#创建环境
mkvirtualenv name #在~/.virtualenvs下安装name的需拟环境
#激活工作环境（workon后不加任何东西可以列出所有虚拟环境）
workon name
#退出当前虚拟环境
deactivate
#删除虚拟环境，需先退出
rmvirtualenv name
#列出所有虚拟环境
lsvirtualenv
#列出所有已安装的包
lssitepackages
#如果想看主环境的已安装包用（在主环境下执行）
dpkg -l | grep nvinfer
2、设置cuda环境变量，解决nvcc -V找不到命令
cd ~
vim .bashrc
末尾添加下面三行代码
export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
export CUDA_ROOT=/usr/local/cuda
激活
source .bashrc
3、更新
sudo apt-get update
4、安装pytorch 和 torchvision
最好对照官网的教程操作，重点留意版本号
https://forums.developer.nvidia.com/t/pytorch-for-jetson-version-1-8-0-now-available/72048
这里我把官方教程搬运过来讲解：
#先把官方的torch包下载下来，这里我是魔法上网下载的，因为实在太慢了，如果换成清华源的话找都找不到。自己下的话要注意版本号，我官方镜像自带的是python3.6.9所以是cp36，还要注意架构是Aarch64，不要下错了。就算下错了也没关系，笑死，根本装不上。
wget https://nvidia.box.com/shared/static/p57jwntv436lfrd78inwl7iml6p13fzh.whl -O torch-1.8.0-cp36-cp36m-linux_aarch64.whl
#安装需要的包
sudo apt-get install python3-pip libopenblas-base libopenmpi-dev
pip3 install Cython
pip3 install numpy
#安装刚才下载的torch包
pip3 install torch-1.8.0-cp36-cp36m-linux_aarch64.whl
验证是否成功
python3
>>>import torch
>>>print(torch.__version__)
#出现版本号就算成功
接下来安装torchvision。下面是对应版本号
#先安装必要的包
sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev
#原谅我这里又用了魔法上网(version是pytorch对应torchvision的版本号，我这里填的是 v0.9.0)
git clone --branch <version> https://github.com/pytorch/vision torchvision
#进入到上一步下载的文件夹中
cd torchvision
#设置临时变量(x就是你的版本号，我这里是9)
export BUILD_VERSION=0.x.0
#开始安装（注意这里--user是安装在主环境，如果是虚拟环境要改为 --你的环境名）
python3 setup.py install --user
验证是否成功
python3
>>>import torchvision
>>>print(torchversion.__version__)
#出现版本号就算成功
5、安装yolov5必须环境
接着安装scipy，首先要安装前面这三样，不然会出错
sudo apt-get install liblapack-dev
sudo apt-get install libblas-dev
sudo apt-get install gfortran
pip3 install scipy
#接着安装剩下的包（超级慢）
pip3 install matplotlib pillow pyyaml tensorboard tqdm
#安装seaborn时会自动安装matplotlib和pandas
pip3 install seaborn
恭喜，完成上面这几部步就可以跑yolov5了，我试过3.0 & 4.0 & 5.0 版本都可以跑。自己下载https://github.com/ultralytics/yolov5.git运行python3 detect.py试一下，这里不详细说明了。
执行上面基本忽略使用魔法的时间就需要3.5个小时，希望你还有耐心听我讲完tensorrt的加速步骤。
二、TensorRt加速
我用的是tensorrtx，附上下载地址：https://github.com/wang-xinyu/tensorrtx.git
你也可以尝试jetson-inference，我试过可以，不过安装过程繁琐，可以参考这位大神的教程：https://blog.csdn.net/qianbin3200896/article/details/108949723
如果你想要jetson-inference，而同时又用了虚拟环境，需要注意建立jetson和虚拟环境的软链接，参考这位大神的5.4的方法三。https://blog.csdn.net/u011119817/article/details/99679350
言归正传，
tensorrtx上有他自己官方的教程，这里我搬运过来：
#1、需要安装pycuda
pip3 install pycuda
#2、将tensorrtx/yolov5/gen_wts.py复制到之前（上面的第5步）你下载的yolov5文件夹中。如果你上一步没有下载，也可以现在下载
git clone https://github.com/ultralytics/yolov5.git
#3、下载官方的权重文件'yolov5s.pt'，也可以用自己的。我用的是yolov5(4.0版本)训练出来的模型。
#4、执行gen_wts.py生成.wts文件。
python gen_wts.py yolov5s.pt
#5、接下来去到目录tensorrtx下的yolov5文件夹
#老规矩，创建一个build文件,并生成生成makeFile
mkdir build
cd build
cmake ..
#6、将yololayer.h里的CLASS_NUM修改成你的。因为官方用的是coco数据集，所以默认是80。
#7、执行makeFile。（每次修改为CLASS_NUM都要make一次）
make
#8、将第4步生成的.wts文件复制到tensorrtx/yolov5里。
#9、生成.engine文件（我用的是yolov5s，所以结尾用s）
sudo ./yolov5 -s ../yolov5s.wts yolov5s.engine s
如果你训练时是自定义depth_multiple 和 width_multiple就这样写：
sudo ./yolov5 -s ../yolov5.wts yolov5.engine c 0.17 0.25
在tensorrtx 5.0里也更新了yolov5的P6模型：
sudo ./yolov5 -s ../yolov5.wts yolov5.engine s6
#10、用他自带的图片(在samples里有两张图片)测试一下
sudo ./yolov5 -d yolov5s.engine ../samples
你也可以用自带的python文件来测试，不过我执行报错
python yolov5_trt.py
三、调用usb摄像头
我写的是c++，如果你用python请参考这位大神https://blog.csdn.net/weixin_45569617/article/details/108145046
在c++中，tensorrtx官方并没有给出调用摄像头的代码，我在github找到一位大神的代码，不过版本比较旧，于是在他的基础上稍作修改。里面我删掉了生成.engine文件的代码，需要你执行完上面的步骤生成了.engine再看下面的代码。
用下面代码替换掉tensorrtx\yolov5\yolov5.cpp
#include (end - start).count() << "ms" << std::endl;
int fps = 1000.0/std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
std::vector<std::vector<Yolo::Detection>> batch_res(fcount);
for (int b = 0; b < fcount; b++) {
auto& res = batch_res[b];
nms(res, &prob[b * OUTPUT_SIZE], CONF_THRESH, NMS_THRESH);
}
for (int b = 0; b < fcount; b++) {
auto& res = batch_res[b];
//std::cout << res.size() << std::endl;
//cv::Mat img = cv::imread(img_dir + "/" + file_names[f - fcount + 1 + b]);
for (size_t j = 0; j < res.size(); j++) {
cv::Rect r = get_rect(frame, res[j].bbox);
cv::rectangle(frame, r, cv::Scalar(0x27, 0xC1, 0x36), 2);
std::string label = my_classes[(int)res[j].class_id];
cv::putText(frame, label, cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2);
std::string jetson_fps = "Jetson Nano FPS: " + std::to_string(fps);
cv::putText(frame, jetson_fps, cv::Point(11,80), cv::FONT_HERSHEY_PLAIN, 3, cv::Scalar(0, 0, 255), 2, cv::LINE_AA);
}
//cv::imwrite("_" + file_names[f - fcount + 1 + b], img);
}
cv::imshow("yolov5",frame);
key = cv::waitKey(1);
if (key == 'q'){
break;
}
fcount = 0;
}
capture.release();
// Release stream and buffers
cudaStreamDestroy(stream);
CUDA_CHECK(cudaFree(buffers[inputIndex]));
CUDA_CHECK(cudaFree(buffers[outputIndex]));
// Destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();
return 0;
}
改完后执行
cd build
make
sudo ./yolov5 -v yolov5s.engine
总结
这些方法并不是我原创的，都是通过测试大神们的方法后总结的，完全没有创新。像在做排除法，类似于培养杂交水稻。