玩转NVIDIA Jetson AGX Xavier(14)--- 在jetson上运行tensorflow+tensorrt的官方demo

本文介绍在jetson Xavier上如何从一个tensorflow的网络搭建到部署的完整过程。

  • 环境准备
    jetson Xavier jetpack4.2 (已安装tensorrt)
    tensorflow安装,请点击
  • demo位置
/usr/src/tensorrt/samples/python/end_to_end_tensorflow_mnist

具体步骤

  1. 打开demo文件位置,有四个文件
    model.py tensorflow框架下lenet的网络模型
    README.md 操作步骤介绍,按照这个直接跑会出问题,按照我下边修改之后的运行
    requirements.txt 需要的python插件,里边包括numpy pycuda tensorflow 等4个插件
    sample.py 加载保存后的模型进行部署的程序

  2. 修改requirements.txt
    这里的修改是把pycuda和tensorflow两个插件删掉,这两个需要单独安装。

  3. 安装没问题的包

    sudo python3 -m pip install -r requirements.txt
    这里会安装numpy 和其他一个库,忘了名字了。

  4. 安装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 (一定要这么安装,不然失败)

  5. 训练模型
    先建一个文件夹

    sudo mkdir models
    sudo python3 model.py
    经过一阵的训练会得到.pb的文件在刚刚新建的models文件夹下

  6. 转换.pb文件为uff文件
    这里有个深坑,官方压根没告诉你转换脚本在哪,我也是找了半天,最后发现是在python下的一个目录里,直接上命令

    sudo python3 /usr/lib/python3.6/dist-packages/uff/bin/convert_to_uff.py models/lenet5.pb
    这一步要在demo的目录下执行,成功的话回在models子目录下得到.uff的模型文件,这个文件就是tensorrt用来做部署的文件了。

  7. 进行推理

    sudo python3 sample.py

成功后会显示输入数字和预测的结果数字。

上一篇:Jetson TX2 安装JetPack3.3教程


下一篇:Jetson Nano开箱配置及Tensorflow安装使用