在Nvidia Jetson Nano上使用TensorRT指南

  • 注意：本指南假设您使用的是Ubuntu 18.04。如果您使用的是Windows，请参考这些说明来设置您的计算机以使用TensorRT。

步骤1：在Ubuntu机器上安装TensorRT

按照这里的说明进行操作。请确保按照tar文件的说明进行操作，除非您之前已经使用.deb文件安装了CUDA。

步骤2：在Jetson Nano上安装TensorRT

  • 设置一些环境变量，以便将nvcc添加到$PATH中。将以下行添加到您的~/.bashrc文件中。
# 将以下内容添加到您的.bashrc文件中
export CUDA_HOME=/usr/local/cuda
# 将CUDA编译器添加到PATH中
export PATH=$CUDA_HOME/bin:$PATH
# 添加库
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
  • 测试对 .bashrc 的更改。
source ~/.bashrc
nvcc --version

您应该会看到类似以下的输出：

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2018 NVIDIA Corporation
Built on ...
Cuda compilation tools, release 10.0, Vxxxxx
  • 切换到您的 virtualenv 并安装 PyCUDA。
# 这需要一些时间。
pip install pycuda
  • 之后，您还需要设置 PYTHONPATH，以便将您的 dist-packages 目录包含在您的 virtualenv 中。将以下内容添加到您的 .bashrc 中。这需要这样做，因为 tensorrt 的 Python 绑定在 dist-packages 中，而该文件夹通常对 virtualenv 不可见。为了使其可见，我们将其添加到 PYTHONPATH 中。
export PYTHONPATH=/usr/lib/python3.6/dist-packages:$PYTHONPATH
  • 切换到您的 virtualenv 并导入 tensorrt，以测试此更改。
> import tensorrt as trt
>  # 这个导入应该成功

第三步：训练、冻结和导出模型为 TensorRT 格式（uff）

在训练完成 linear 模型后，您将得到一个带有 .h5 扩展名的文件。

# 您将在 models 文件夹中得到一个 Linear.h5 文件
python manage.py train --model=./models/Linear.h5 --tub=./data/tub_1_19-06-29,...

# （可选）将 './models/Linear.h5' 从您的台式电脑复制到 Jetson Nano 上的工作目录（~mycar/models/）

# 使用 donkeycar/scripts 中的 freeze_model.py 冻结模型；冻结的模型以协议缓冲区的形式保存。
# 此命令还会导出模型的一些元数据，保存在 ./models/Linear.metadata 中
python ~/projects/donkeycar/scripts/freeze_model.py --model=~/mycar/models/Linear.h5 --output=~/mycar/models/Linear.pb

# 将冻结的模型转换为 UFF 格式。下面的命令将创建一个名为 ./models/Linear.uff 的文件
cd /usr/lib/python3.6/dist-packages/uff/bin/
python convert_to_uff.py ~/mycar/models/Linear.pb

现在将转换后的 uff 模型和 metadata 复制到 Jetson Nano 上。

第四步

  • 在 myconfig.py 中选择模型类型为 tensorrt_linear。
DEFAULT_MODEL_TYPE = `tensorrt_linear`
  • 最后，您可以执行以下命令
# 在将 `uff` 模型通过 scp 复制到 Nano 后
python manage.py drive --model=./models/Linear.uff