本文介绍在jetson Xavier上如何从一个tensorflow的网络搭建到部署的完整过程。
- 环境准备
jetson Xavier jetpack4.2 (已安装tensorrt)
tensorflow安装,请点击 - demo位置
/usr/src/tensorrt/samples/python/end_to_end_tensorflow_mnist
具体步骤
-
打开demo文件位置,有四个文件
model.py tensorflow框架下lenet的网络模型
README.md 操作步骤介绍,按照这个直接跑会出问题,按照我下边修改之后的运行
requirements.txt 需要的python插件,里边包括numpy pycuda tensorflow 等4个插件
sample.py 加载保存后的模型进行部署的程序 -
修改requirements.txt
这里的修改是把pycuda和tensorflow两个插件删掉,这两个需要单独安装。 -
安装没问题的包
sudo python3 -m pip install -r requirements.txt
这里会安装numpy 和其他一个库,忘了名字了。 -
安装pycuda
这一步把我坑坏了,安装了一天都没成功,Google了各种论坛,最后找到原因是cuda的环境变量配置的问题。直接上最后的安装方案sudo gedit ~/.bashrc
export PATH=/usr/local/cuda-10.0/bin:$PATH
source ~/.bashrc
export LD_LIBRARY_PATH=/usr/local/cuda-10.0/lib64
pip3 install pycuda --user (一定要这么安装,不然失败) -
训练模型
先建一个文件夹sudo mkdir models
sudo python3 model.py
经过一阵的训练会得到.pb的文件在刚刚新建的models文件夹下 -
转换.pb文件为uff文件
这里有个深坑,官方压根没告诉你转换脚本在哪,我也是找了半天,最后发现是在python下的一个目录里,直接上命令sudo python3 /usr/lib/python3.6/dist-packages/uff/bin/convert_to_uff.py models/lenet5.pb
这一步要在demo的目录下执行,成功的话回在models子目录下得到.uff的模型文件,这个文件就是tensorrt用来做部署的文件了。 -
进行推理
sudo python3 sample.py
成功后会显示输入数字和预测的结果数字。