PyTorch build on Jetson Nano

PyTorch

この記事ではPyTorchをJetson Nanoでビルドする方法について解説します。現時点でJetson Linuxがサポートするpythonのバージョンが3.6であるため、python3.7以上のバージョンでPyTorchを使用する場合はJetson公式のPyTorchライブラリが存在しないため自分でPyTrochをビルドする必要があります。

この記事ではpython3.8の例を説明しますがそれ以外のバージョンでもほぼ同じ仕組みで動作すると思います。

install packages

python3.8とPyTorchで必要なパッケージをインストールします。

sudo apt-get update sudo apt-get install -y \ python3.8 python3.8-dev \ ninja-build git cmake clang \ libopenmpi-dev libomp-dev ccache \ libopenblas-dev libblas-dev libeigen3-dev \ python3-pip libjpeg-dev \ gnupg2 curl
Code language: Bash (bash)

install cmake

標準のcmakeはPyTorchをビルドするには古いため、PyTorchビルドに対応するcmakeをインストールします。

# cmake sudo apt-get install -y software-properties-common lsb-release wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | gpg --dearmor - | sudo tee /etc/apt/trusted.gpg.d/kitware.gpg >/dev/null sudo apt-add-repository "deb https://apt.kitware.com/ubuntu/ $(lsb_release -cs) main" sudo apt-get update sudo apt-get install -y cmake
Code language: Bash (bash)

install CUDA and cuDNN

CUDAとcuDNNをインストールします。

# cuda and cudnn sudo apt-get install -y nvidia-cuda nvidia-cudnn8
Code language: Bash (bash)

remove package

PytorchとTorchVisionをビルドする前にaptでインストールされているnumpyがコンフリクトするため削除します。

sudo apt remove -y python3-numpy
Code language: Bash (bash)

build and install PyTorch with python3.8

python3.8 -m pip install -U pip python3.8 -m pip install -U setuptools python3.8 -m pip install -U wheel mock pillow python3.8 -m pip install scikit-build python3.8 -m pip install cython Pillow ## download PyTorch v1.11.0 with all its libraries git clone -b v1.11.0 --depth 1 --recursive --recurse-submodules --shallow-submodules https://github.com/pytorch/pytorch.git ( cd pytorch python3.8 -m pip install -r requirements.txt wget https://raw.githubusercontent.com/otamajakusi/build_jetson_nano_libraries/main/pytorch/pytorch-1.11-jetson.patch patch -p1 < pytorch-1.11-jetson.patch )
Code language: Bash (bash)

python3.8のパッケージをインストールし、PyTorchをクローンしパッチを当てます。パッチの詳細はこちらをご覧ください。

以下はpath/to/pytorch-1.11-jetson.patchの内容です。こちらからダウンロードできます。

diff --git a/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h b/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h index 0327868..e484fba 100644 --- a/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h +++ b/aten/src/ATen/cpu/vec/vec256/vec256_float_neon.h @@ -26,6 +26,7 @@ inline namespace CPU_CAPABILITY { // https://bugs.llvm.org/show_bug.cgi?id=45824 // Most likely we will do aarch32 support with inline asm. #if defined(__aarch64__) +#if defined(__clang__) || (__GNUC__ > 8 || (__GNUC__ == 8 && __GNUC_MINOR__ > 3)) #ifdef __BIG_ENDIAN__ #error "Big endian is not supported." @@ -715,5 +716,6 @@ Vectorized<float> inline fmadd(const Vectorized<float>& a, const Vectorized<floa } #endif /* defined(aarch64) */ +#endif /* defined(__clang__) */ }}} diff --git a/aten/src/ATen/cuda/CUDAContext.cpp b/aten/src/ATen/cuda/CUDAContext.cpp index 1751128..a090e70 100644 --- a/aten/src/ATen/cuda/CUDAContext.cpp +++ b/aten/src/ATen/cuda/CUDAContext.cpp @@ -24,6 +24,7 @@ void initCUDAContextVectors() { void initDeviceProperty(DeviceIndex device_index) { cudaDeviceProp device_prop; AT_CUDA_CHECK(cudaGetDeviceProperties(&device_prop, device_index)); + device_prop.maxThreadsPerBlock = device_prop.maxThreadsPerBlock / 2; device_properties[device_index] = device_prop; } diff --git a/aten/src/ATen/cuda/detail/KernelUtils.h b/aten/src/ATen/cuda/detail/KernelUtils.h index b36e78c..dea597f 100644 --- a/aten/src/ATen/cuda/detail/KernelUtils.h +++ b/aten/src/ATen/cuda/detail/KernelUtils.h @@ -19,7 +19,7 @@ namespace at { namespace cuda { namespace detail { // Use 1024 threads per block, which requires cuda sm_2x or above -constexpr int CUDA_NUM_THREADS = 1024; +constexpr int CUDA_NUM_THREADS = 512; // CUDA: number of blocks for threads. inline int GET_BLOCKS(const int64_t N, const int64_t max_threads_per_block=CUDA_NUM_THREADS) {
Code language: Diff (diff)

PyTorchをビルドします。Jetson Nanoのメモリが少ないためMAX_JOBSを2としています。MAX_JOBSを増やすと速くビルドできますが非常に高い確率でハングアップします。

export BUILD_CAFFE2_OPS=OFF export USE_FBGEMM=OFF export USE_FAKELOWP=OFF export BUILD_TEST=OFF export USE_MKLDNN=OFF export USE_NNPACK=OFF export USE_XNNPACK=OFF export USE_QNNPACK=OFF export USE_PYTORCH_QNNPACK=OFF export USE_CUDA=ON export USE_CUDNN=ON export TORCH_CUDA_ARCH_LIST="5.3;6.2;7.2" export USE_NCCL=OFF export USE_SYSTEM_NCCL=OFF export USE_OPENCV=OFF export MAX_JOBS=2 # set path to ccache export PATH=/usr/lib/ccache:/usr/local/cuda/bin:$PATH # set clang compiler export CC=clang export CXX=clang++ # create symlink to cublas # ln -s /usr/lib/aarch64-linux-gnu/libcublas.so /usr/local/cuda/lib64/libcublas.so # start the build ( cd pytorch python3.8 setup.py bdist_wheel )
Code language: Bash (bash)

このビルド時間はおおよそ13時間かかります。

インストールします。

find pytorch/dist -type f|xargs python3.8 -m pip install
Code language: Bash (bash)

build and install TorchVision with python3.8

TorchVisionをクローンしビルドします。

# torch vision git clone --depth=1 https://github.com/pytorch/vision torchvision -b v0.12.0 ( cd torchvision export TORCH_CUDA_ARCH_LIST='5.3;6.2;7.2' export FORCE_CUDA=1 export MAX_JOBS=2 python3.8 setup.py bdist_wheel )
Code language: Bash (bash)

このビルド時間はおおよそ16分かかります。

インストールします。

find torchvision/dist -type f|xargs python3.8 -m pip install
Code language: Bash (bash)

以上です。

参考記事