Use PyTorch upstream image for Jetson

.github/workflows/docker-build-jetson.yaml

@@ -49,7 +49,7 @@ jobs:
           tags: ${{ steps.meta.outputs.tags }}
           platforms: linux/arm64
           labels: ${{ steps.meta.outputs.labels }}
-          build-args: upstream_image=nvcr.io/nvidia/l4t-cuda:11.4.19-devel
+          build-args: upstream_image=nvcr.io/nvidia/l4t-pytorch:r35.1.0-pth1.13-py3
       -
         name: Image digest
         run: echo ${{ steps.docker_build.outputs.digest }}

build.bash

@@ -20,8 +20,11 @@
 
 set -eo pipefail
 
+# Upstream images for x86_64 and aarch64
+# For x86_64, use the CUDA 11.4.3 image
+# For aarch64, use the L4T PyTorch image
 UPSTREAM_X86_64="nvcr.io/nvidia/cuda:11.4.3-devel-ubuntu20.04"
-UPSTREAM_AARCH64="nvcr.io/nvidia/l4t-cuda:11.4.19-devel"
+UPSTREAM_AARCH64="nvcr.io/nvidia/l4t-pytorch:r35.1.0-pth1.13-py3"
 
 # Check that the current user has UID 1000.
 if [ $(id -u) -ne 1000 ]

dcist-master/Dockerfile

@@ -1,3 +1,5 @@
+# The upstream image depends on the architecture, and it is defined when
+# building the image. Please see build.bash for more details
 ARG upstream_image
 FROM $upstream_image
 
@@ -71,8 +73,9 @@ USER $USER
 
 # Install ROS Noetic
 RUN sudo sh -c 'echo "deb http://packages.ros.org/ros/ubuntu $(lsb_release -sc) main" > /etc/apt/sources.list.d/ros-latest.list' \ && sudo /bin/sh -c 'curl -s https://raw.githubusercontent.com/ros/rosdistro/master/ros.asc | apt-key add -' \
- && sudo apt-get update \
- && sudo apt-get install -y \
+ && sudo apt-get update
+
+RUN sudo apt-get install -y \
     python3-catkin-tools \
     python3-rosdep \
     python3-rosinstall \
@@ -103,8 +106,8 @@ RUN sudo apt-get install -y \
 # Install geographiclib_datasets required by mavros
 RUN cd /opt/ros/noetic/lib/mavros/ && sudo ./install_geographiclib_datasets.sh
 
-# Install pytorch
-RUN sudo sh -c 'pip3 install torch==1.8.2+cu111 torchvision==0.9.2+cu111 torchaudio==0.8.2 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html'
+# Install pytorch only if it is not installed (avoid re-installing in Jetson)
+RUN sudo sh -c 'pip3 freeze | grep torch || pip3 install torch==1.8.2+cu111 torchvision==0.9.2+cu111 torchaudio==0.8.2 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html'
 
 # Install dependencies for MOCHA
 RUN sudo sh -c 'pip3 install zmq rospkg defusedxml utm'