一拿到熱騰騰的 Jetson Nano，當然第一步就是立刻裝系統，上環境。但是在裝 TensorFlow 和 PyTorch 可是栽了很多坑。因為 Jetson Nano 是 arm64 的 SoC，雖然 TensorFlow 和 PyTorch CPU only 有 for arm64 的 wheel 可以安裝，但 GPU support 的卻都沒有！都用上 Jetson Nano 怎麼可以不用 GPU 呢？
經過一段時間的試誤和 google 後，有以下結論
- 不要自己 build
因為太麻煩了！要裝很多的依賴， CPU 不夠強要 build 超久，還有蠻大的機率 build 失敗！想玩個 AI 搞這麼痛苦幹嘛呢！ - 用 nvidia 官方 official build 好的來裝
nvidia 真是太佛心了，該有的都 build 好，直接照他指令一個個下就好，只是不知為何官網都沒有寫，反而要搜尋到官方論壇才有。 - Python 要使用 3.6 版
因為 JetPack 自帶的 python 是 3.6 ，所以官方 build 好的也都只有 for 3.6，如果自己安裝新版的 python 就不能使用 - 可以使用官方提供的 docker image 或是 Dockerfile
這個蠻不錯的，如果開發好的應用要放進 docker 裡跑，直接使用官方的 image 當 base 下去做。另外如果需要在 docker 裡面安裝 openCV 或是 ROS （這兩個也是不好搞），官方也有對應的 Dockerfile 範例可以拿來改
好！總之我們要站在巨人的肩膀上，不要自己造輪子！首先讓我們看看如何使用 official build 的方式安裝
使用 nvidia official build 安裝
安裝 TensorFlow 2
詳細可參考 Official TensorFlow for Jetson Nano! 這邊簡單列出需要的指令
# 準備依賴
sudo apt-get install libhdf5-serial-dev hdf5-tools libhdf5-dev zlib1g-dev zip libjpeg8-dev liblapack-dev libblas-dev gfortran
sudo apt-get install python3-pip
sudo pip3 install -U pip testresources setuptools==49.6.0
sudo pip3 install -U numpy==1.16.1 future==0.18.2 mock==3.0.5 h5py==2.10.0 keras_preprocessing==1.1.1 keras_applications==1.0.8 gast==0.2.2 futures protobuf pybind11
# 開始安裝 TF-2.x
sudo pip3 install --pre --extra-index-url https://developer.download.nvidia.com/compute/redist/jp/v45 tensorflow安裝完後，因為 2.x 有新增一些指令，要給定 cpu type 才不會發生 Illegal instruction (core dumped) 的錯誤
# 把這行加到 .bashrc
export OPENBLAS_CORETYPE=ARMV8裝完後驗證一下是否有使用 GPU
import tensorflow as tf
if tf.test.gpu_device_name():
print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))
else:
print("Please install GPU version of TF")安裝 PyTorch
一樣詳細可參考 PyTorch for Jetson – version 1.8.0 now available，這邊簡單列出需要的指令
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-1.8.0-cp36-cp36m-linux_aarch64.whl順便也安裝一下 torchvision，但這個要 build 所以會花一段不算短的時間。這邊以 0.9.0 為例
sudo apt-get install libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev
git clone --branch v0.9.0 https://github.com/pytorch/vision torchvision
cd torchvision
export BUILD_VERSION=0.9.0
python3 setup.py install --user此時如果跑 pytorch 可能會出現錯誤，給定 LD_PRELOAD 即可解
ImportError: /usr/lib/aarch64-linux-gnu/libgomp.so.1: cannot allocate memory in static TLS block
# 一樣把這行加到 .bashrc
export LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libgomp.so.1裝完後驗證一下是否有使用 GPU
import torch
# True 代表成功
torch.cuda.is_available()
# 以下只是取得 GPU 資訊
torch.cuda.current_device()
torch.cuda.device(0)
torch.cuda.device_count()
torch.cuda.get_device_name(0)使用 Docker image or Dockerfile
官方在 github 有列出 image name 可以直接 pull，有需要還可以使用 git 裡的 Dockerfile 來修改
其中特別注意的是，要在 docker command 給定 runtime 或是直接到 /etc/docker/daemon.json 設定 default runtime ，這樣在 container 裡面才有 GPU 加速！
# 舉例，在 run 指令後面給定 runtime
docker run -it --runtime=nvidia nvcr.io/nvidia/l4t-pytorch:r32.5.0-pth1.7-py3 bash// 更方便並防忘記的方法，就是直接修改 /etc/docker/daemon.json
{
"runtimes": {
"nvidia": {
"path": "nvidia-container-runtime",
"runtimeArgs": []
}
},
"default-runtime": "nvidia" // 加這一行
}That’s it! 希望可以讓大家少走點彎路，趕快開始玩 AI！
延伸閱讀：如何在 Jetson Nano 上安裝 OpenCV？
參考資料
Official TensorFlow for Jetson Nano!
PyTorch for Jetson – version 1.8.0 now available
https://github.com/dusty-nv/jetson-containers
How to check if pytorch is using the GPU?
Check whether Tensorflow is running on GPU
Illegal instruction(core dumped) error on Jetson Nano
