在NVIDIA Jetson TX2上源码编译并使用Pytorch

前言

TX2作为一个嵌入式平台的深度学习端，具备不错的GPU性能，我们可以发现TX2的GPU的计算能力是6.2。这意味着TX2对半精度运算有着良好的支持，我们完全可以在桌面端训练好模型，然后移植到TX2上利用半精度运行进行推理，这样可以达到生产落地的效果。

就算本篇文章主要展示了如何在TX2中源码编译Pytorch-1.0。

首先我们需要一个相对纯净的jetpack系统，3.2-3.3版本(最新的4.1.1也可以)都可以，所以我们最好将TX2的系统重新刷一遍，以免造成一些其他不兼容的错误。

刷系统：从NVIDIA官网下载TX2的系统包：https://developer.nvidia.com/embedded/jetpack

步骤

接下来我们严格按照步骤来进行Pytorch的源码安装。

在JetPack-3.2系统中，我们的python一般有两个版本，python命令对应着python2.7。而python3命令对应着python3.5。这里我们使用python3作为编译环境，大家一定要分清楚这两个不同python版本的命令集，否则会造成编译错误。

可以通过which python3来查看当前的python3.5的执行环境。

依赖件

首先安装依赖件：

注意我们使用的命令pip3对应着python3，如果你不清楚系统中的pip和python的绑定信息，采用pip --version查看当前命令在哪个python中，例如我这边则是：

pip3 --version
pip3 9.0.1 from path/to/lib/python3.5/site-packages/pip (python 3.5)

总之我们要使用对应python3的pip命令，首先安装pip3，然后在python3环境中安装一些必要组件。

sudo apt install libopenblas-dev libatlas-dev liblapack-dev
sudo apt install liblapacke-dev checkinstall # For OpenCV
sudo apt-get install python3-pip

pip3 install --upgrade pip3==9.0.1
sudo apt-get install python3-dev

sudo pip3 install numpy scipy # 这个花费的时间略长，20分钟到30分钟左右
sudo pip3 install pyyaml
sudo pip3 install scikit-build
sudo apt-get -y install cmake
sudo apt install libffi-dev
sudo pip3 install cffi

安装完之后，我们添加cudnn的lib和include路径，为什么要执行这一步，因为我们在刷好机后，cuda和cudnn也已经安装好，但是JetPack系统中的cudnn路径和我们一般ubuntu系统中的路径略有不同(为什么不同看这里：https://oldpan.me/archives/pytorch-gpu-ubuntu-nvidia-cuda90)，这时需要我们将cudnn的路径添加到环境变量中并激活：

sudo gedit ~/.bashrc
export CUDNN_LIB_DIR=/usr/lib/aarch64-linux-gnu
export CUDNN_INCLUDE_DIR=/usr/include
source ~/.bashrc

下载Pytorch源码包

我们从github上直接拷贝最新的Pytorch源码包，然后利用pip3安装好所有必备的库，并对第三方库进行更新。

git clone http://github.com/pytorch/pytorch
cd pytorch
sudo pip3 install -U setuptools
sudo pip3 install -r requirements.txt
git submodule update --init --recursive

编译

在编译之前，我们先开启TX2的最大功率模式，这样可以使我们的编译速度稍微快一些：

sudo nvpmodel -m 0         # 切换工作模式到最大
sudo  ~/jetson_clocks.sh   # 强制开启风扇最大转速 

这样风扇就开始呼呼转了，我们开始编译吧。

进入Pytorch源码目录后，我们首先执行下面这一句首先编译Pytorch的开发组件：

python3 setup.py build_deps

漫长的编译后(2小时)，我们继续执行以下的命令：

NO_SYSTEM_NCCL=1 DEBUG=1 sudo python3 setup.py build develop

同样是漫长的编译(2小时)，等待后不出意外就编译成功了！

后续操作

后续我们继续安装Pytorch推荐搭配的组件torchvision。

sudo apt clean
sudo apt-get install libjpeg-dev zlib1g-dev

git clone https://github.com/python-pillow/Pillow.git
cd Pillow/
sudo python3 setup.py install
sudo pip3 install pandas # ~20-30 min
sudo pip3 install Cython
sudo pip3 install scikit-image

sudo pip3 --no-cache-dir install torchvision

测试Pytorch是否安装成功

编译完成之后，如果安装成功，会出现下面的信息：

需要注意的点

在JetPack系统中，我们不论是执行编译还是执行代码，都必须加上sudo，否则会出现编译失败或者cuda error的情况。