Jetson Nanoへのtorchaudioインストール

summary

ここでtorch1.10, torchaudio0.10, torchvision0.11のwheelを公開しています。torch, torchvisionのwheelはここから拝借しました。

テスト環境は次の通りです。

• Jetson Nano 2GB / 4GB
• Jetpack 4.6.0 (Python 3.6)

背景

Jetson Nanoでtorchaudioを使いたかったのですが、公式にはdockerを使う方法しか乗っていません。wheelがどこかに転がっていないか結構探したけど見つかりませんでした。
おとなしくdockerを使え、という話でもありますが、自分でビルドしてみました。ちなみに、CUDAを使わない場合はここのwheelを使えます。

ビルド手順

Jetson用PytorchのDockerfileや、JetsonNano用のtorch/torchvisionwheelを公開しているサイトの手順を参考にしました。
aptで入れたcmakeを使おうとするとバージョンが古すぎる、と怒られるので、まずcmakeを新しくします。2022/4/7現在、最新版の3.23.0にするとtorchaudioがmakeできないので、少し前のやつ(3.22.2等)を使います。

wget https://github.com/Kitware/CMake/releases/download/v3.22.2/cmake-3.22.2.tar.gz
tar zxvf cmake-3.22.2.tar.gz
cd cmake-3.22.2
./bootstrap
make
sudo make install

最後に、torchaudioをクローンしてビルドします。

git clone -b release/0.10 https://github.com/pytorch/audio.git torchaudio
cd torchaudio
sudo apt-get install clang
sudo su
export USE_FBGEMM=OFF
export BUILD_CAFFE2_OPS=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=4
export PATH=/usr/lib/ccache:$PATH
# set clang compiler
export CC=clang
export CXX=clang++
sudo ln -s /usr/lib/aarch64-linux-gnu/libcublas.so /usr/local/cuda/lib64/libcublas.so
export CUDACXX=/usr/local/cuda-10.2/bin/nvcc
python3 setup.py clean
python3 setup.py bdist_wheel

dist/以下にwheelができているはずです！
ちなみに、私ははじめ手元にあったJetsonNano2GBでビルドを試みたのですが、RAMをすぐさま使い切った後フリーズしてしまったので4GBに変えて作業しました。