Installing Torch and Torch_Scatter with CUDA Support on AArch64 Devices: A Guide

1. Environment

Device: Nvidia Jetson AGX Orin / Jetson Orin N / Drive Orin / ....
Ubuntu: 20.04
Architecture: aarch64(arm64)

2. Question

If we need to install torch and torch_scatter with CUDA support on aarch64 devices, we can install torch according to this guide (3.1). However, there is no precomplied installation package -- torch_scatter, we need to install it from the source code.

3. Solution

1. Installing `torch` (CUDA, aarch64)

https://docs.nvidia.com/deeplearning/frameworks/install-pytorch-jetson-platform/index.html#abstract

2. Installing `torch-scatter`

Installing torch-scatter with CUDA support from source on ARM-based devices, like the NVIDIA Jetson series, can be more complex due to the need to ensure the CUDA environment is correctly configured and may require specific compilation adjustments for ARM architecture. Here are the general steps for the installation process:

Prerequisites

Ensure CUDA is Installed: Before starting, make sure that the CUDA Toolkit is installed on your device. You can verify the installation and version of CUDA by running nvcc --version.
Install PyTorch: Ensure that a CUDA-supported version of PyTorch is installed for your device. For NVIDIA Jetson devices, you might need to use the PyTorch pre-built packages optimized for ARM architecture provided by NVIDIA, or compile PyTorch from source.
Install Necessary Compilation Tools: Make sure your system has the necessary compilation tools installed, such as gcc, g++, and cmake.

Installation Steps

Clone the torch-scatter Repository

git clone https://github.com/rusty1s/pytorch_scatter.git
cd pytorch_scatter

Set CUDA Environment Variables (if necessary)

Ensure the environment variables are correctly set to point to the CUDA toolchain. You might need to set the CUDA_HOME environment variable:

export CUDA_HOME=/usr/local/cuda

Additionally, make sure the CUDA-related directories (like bin and lib64) are added to PATH and LD_LIBRARY_PATH:

export PATH=/usr/local/cuda/bin:${PATH}
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH}

Install torch-scatter

Use pip to install from source, it takes a few minutes, please wait patiently.

pip install .

This command will initiate the compilation process, which may take some time depending on your system configuration. The compilation will consider CUDA support.

Considerations

Check Compatibility: Ensure the version of torch-scatter is compatible with your installed PyTorch and CUDA versions. You might need to consult the documentation or GitHub repository issues of torch-scatter for information on version compatibility.

Given that Jetson devices and other ARM-based systems might have specific requirements and constraints, the steps outlined above might need to be adjusted for your specific situation. Be prepared to debug and troubleshoot based on the specific issues encountered during the installation process.

rusty1s / pytorch_scatter