Docker container example that can use CUDA and Pytorch 2.5.0 with the host running Jetpack 6.1

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.

Hi,

Thanks a lot for sharing.
Here are some notes which might also be helpful.

  1. The PyTorch container for Jetson can be found in the below link (tag with iGPU).
    PyTorch | NVIDIA NGC

  2. PyTorch, TorchVision, and TorchAudio prebuilt packages for JetPack 6.1 can be found in the below link:
    https://pypi.jetson-ai-lab.dev/jp6/cu126

Thanks.

1 Like

Thanks for this!

Maybe my assumption is wrong that, I can use any version of these regardless of what version of Jetpack my host is running? Or is this the cause of all my headaches?

Hi I tried to make this work and thanks for posting but I get an unauthorised or unavailable error when I try and access nvcr.io/nvidia/l4t-jetpack:r36.4.0? what am i doing wrong? thanks.

Hi,

Each NGC container might have different requirements or constraints.
Our suggestion is to check the support matrix document for the details.

For example, the latest pytorch:24.12-py3-igpu container:

  • PyTorch version: 2.6.0a0+df5bbc0
  • JetPack version: 6.1

It means the container depends on the JetPack 6.1 software.
Thanks.

Hi, @nav-intel

Have you run the docker command with sudo?
We can pull the image with the command below:

$ sudo docker pull nvcr.io/nvidia/l4t-jetpack:r36.4.0
r36.4.0: Pulling from nvidia/l4t-jetpack
3e5db86eb9ec: Pulling fs layer 
c6a3b84cb8a7: Pulling fs layer 
abcaac0f407d: Pulling fs layer 
8c3aa7bdd0ee: Waiting 
1c034b799890: Waiting 
4a709f97643e: Waiting 
...

Thanks.

Happy Christmas @AastaLLL, (if you celebrate it)

it is downloading now, thanks, what I hope is that I can extract Pytorch. Pytorch audio and Pytorch Vision to install into another docker container build process. I am still a bit of a Docker newbie so getting lots of interesting errors :(

Thanks for your help,

info@nav-intel.ai

I think you need to setup your api key.

First go to: https://org.ngc.nvidia.com/setup/api-key
Add the key. It will give you the exact username and password example. The key becomes your password.

Then run:

    $ docker login nvcr.io

It gives you 3 choices. Pick the first one.
Then…

Username:
        $oauthtoken
Password:
        The key it wants.

Now you can download the nvcr.io dockers!

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.