I’ll later be doing an actual guide that gets you from flashing Jetpack 61 to getting cuda installed then docker so all the commands are in one place.
I’m sure this is easy for most of you, but I’m very happy to finally get this running, because I have yet to find a working pytorch container that works on the jetson. If you’ve been struggling, struggle no more.
Some requirements for the host system:
$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Tue_Oct_29_23:53:06_PDT_2024
Cuda compilation tools, release 12.6, V12.6.85
Build cuda_12.6.r12.6/compiler.35059454_0
Dockerfile:
FROM nvcr.io/nvidia/l4t-jetpack:r36.4.0
# Set NVIDIA driver capabilities
ARG NVIDIA_DRIVER_CAPABILITIES
ENV NVIDIA_DRIVER_CAPABILITIES=${NVIDIA_DRIVER_CAPABILITIES:-compute}
# Install required packages and dependencies
RUN apt-get update && apt-get install -y \
wget \
python3-pip \
&& apt-get clean && rm -rf /var/lib/apt/lists/*
RUN pip install numpy==1.24.3
# Fix for cuSPARSELt installation
RUN wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/arm64/cuda-keyring_1.1-1_all.deb
RUN dpkg -i cuda-keyring_1.1-1_all.deb
RUN apt-get update
RUN apt-get -y install libcusparselt0 libcusparselt-dev
RUN apt-get update && apt-get install -y git && apt-get clean && rm -rf /var/lib/apt/lists/*
RUN apt-get update && apt-get install -y libopenblas-dev && rm -rf /var/lib/apt/lists/*
#pytorch
ARG TORCH_INSTALL=https://developer.download.nvidia.com/compute/redist/jp/v61/pytorch/torch-2.5.0a0+872d972e41.nv24.08.17622132-cp310-cp310-linux_aarch64.whl
RUN echo "TORCH_INSTALL=$TORCH_INSTALL"
RUN python3 -m pip install --no-cache-dir "$TORCH_INSTALL"
# Install torchvision
ARG TORCHVISION_VERSION=v0.20.0
ARG PILLOW_VERSION=pillow>=9.1.0
ENV TORCH_CUDA_ARCH_LIST="5.3;6.2;7.2;8.7"
# Echo the branch being installed
RUN echo "Installing TorchVision from branch ${TORCHVISION_VERSION}"
# Clone the TorchVision repository
RUN git clone -b ${TORCHVISION_VERSION} --depth=1 https://github.com/pytorch/vision torchvision
# Change directory into TorchVision and install it
RUN cd torchvision && python3 setup.py install
# Clean up the TorchVision repository
RUN rm -rf torchvision
# Install Pillow
RUN pip3 install --no-cache-dir "${PILLOW_VERSION}"
RUN git clone -b v0.20.0 --depth=1 https://github.com/pytorch/vision torchvision && \
ls torchvision && \
cd torchvision && \
python3 setup.py install
# Ensure runtime NVIDIA capabilities are available
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_REQUIRE_CUDA="cuda>=11.0"
ENV PATH="/usr/local/cuda/bin:${PATH}"
ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"
# Set entrypoint for easier debugging
CMD ["bash"]
Build this bad boy (takes about 400GB):
docker build -t jetpack61-pytorch-image-test .
then when you run it, this is important. You must add -e NVIDIA_DRIVER_CAPABILITIES=compute or it crashes for some reason.
- docker run --rm -it --runtime=nvidia -e NVIDIA_DRIVER_CAPABILITIES=compute test-jetpack61-pytorch-image-test bash
Feel free to ask me to make changes. This is a work in progress and hopefully can make this more slim.