Install pytorch3d in aarch64/arm Linux architecture

[Sar-thak-3]

❓ Questions on how to use PyTorch3D

Is pytorch3d work on aarch64/arm architecture? I want to install this library for my project in this system architecture, but when I try it shows the error of

ERROR: Could not find a version that satisfies the requirement pytorch3d (from versions: none)
ERROR: No matching distribution found for pytorch3d

[bottler]

It may work to build from source. There are no binaries provided, hence the error. And I guess it will be CPU-only.

[Sar-thak-3]

Hi @bottler , I tried setting up repo from source with these commands

git clone https://github.com/facebookresearch/pytorch3d.git
cd pytorch3d && pip install -e .

But Running setup.py develop for pytorch3d keeps running to infinite

[yaojin17]

Same problem. Running setup.py develop for pytorch3d seems running endlessly

[Hydway]

Hi @Sar-thak-3 @yaojin17 , I successfully compiled pytorch3d in a container on arm/v8 and run it in k8s cluster, and I think it also works on host machine. Here is the key steps FYI

FROM jupyter/tensorflow-notebook:aarch64-lab-4.0.3
RUN mkdir -p ~/Downloads
RUN cd ~/Downloads
# https://github.com/conda-forge/miniforge/releases/download/24.3.0-0/Mambaforge-24.3.0-0-Linux-aarch64.sh
COPY Mambaforge-24.3.0-0-Linux-aarch64.sh /root/conda.sh
USER root
RUN chmod +x /root/conda.sh && \
    bash /root/conda.sh -b -p /root/conda
RUN /root/conda/bin/conda init
RUN source ~/.bashrc
ENV PATH /root/conda/bin:$PATH
RUN apt install -y vim-tiny
RUN conda create -y -n pytorch3d python=3.10
RUN echo "conda activate pytorch3d" >> ~/.bashrc

RUN CONDA_OVERRIDE_CUDA="12.0" conda install --name pytorch3d --yes pytorch=2.3.1=cuda120_py310h377a36d_200 -c conda-forge -c nvidia

RUN conda install --name pytorch3d --yes -c conda-forge -c fvcore fvcore

RUN /bin/bash -c ". activate pytorch3d && \
    pip install --upgrade pip && \
    pip install \
    scikit-image \
    matplotlib \
    imageio \
    black \
    isort \
    flake8 \
    flake8-bugbear \
    flake8-comprehensions"

# download the pytorch3d v0.7.6 source code from repo
COPY pytorch3d /root/pytorch3d/
# download the torchvision 0.18.1 source code from repo
COPY vision-0.18.1/vision-0.18.1 /root/torchvision/

RUN /bin/bash -c ". activate pytorch3d && \
    pip install ninja && \
    cd /root/torchvision && \
    pip install . "

WORKDIR /root/pytorch3d
# https://developer.nvidia.com/cuda-12-0-0-download-archive
COPY cuda_12.0.0_525.60.13_linux_sbsa.run /root/cuda_12.0.0_525.60.aarch64.run

RUN sh /root/cuda_12.0.0_525.60.aarch64.run --toolkit  --silent
ENV CUDA_HOME /usr/local/cuda-12.0

RUN /bin/bash -c ". activate pytorch3d && \
    cd /root/pytorch3d && \
    FORCE_CUDA=1 pip install . "

as you can see Pytorch3d needs both arm version of torch and torchvision, and torchvision needs to be compiled manually. And ninja is also required(to speed up compilation process, as you mentioned the compilation seems infinite). When you finally compile the pytorch3d, make sure using FORCE_CUDA=1 to force compile and install the cuda version of pytorch3d. GLHF!