CUDA on the Jetson TX2

Dockerfile

FROM arm64v8/ubuntu:16.04

RUN apt-get update && apt-get install -y --no-install-recommends \
  wget \
  bzip2 \
  && rm -rf /var/lib/apt/lists/*

WORKDIR /tmp

# CUDA
RUN wget http://developer.download.nvidia.com/devzone/devcenter/mobile/jetpack_l4t/013/linux-x64/cuda-repo-l4t-8-0-local_8.0.84-1_arm64.deb
RUN wget http://developer.download.nvidia.com/devzone/devcenter/mobile/jetpack_l4t/013/linux-x64/libcudnn6_6.0.21-1+cuda8.0_arm64.deb
RUN dpkg -R --install *.deb
RUN apt-get update && apt-get install -y --no-install-recommends \
  cuda-toolkit-8.0 \
  libcudnn6 \
  && rm -rf /var/lib/apt/lists/*

# Tegra drivers and tools
RUN wget http://developer.download.nvidia.com/devzone/devcenter/mobile/jetpack_l4t/013/linux-x64/Tegra210_Linux_R28.1.0_aarch64.tbz2 && tar -xvf Tegra210_Linux_R28.1.0_aarch64.tbz2
RUN tar -C / -xvf Linux_for_Tegra/nv_tegra/nvidia_drivers.tbz2 usr/lib/aarch64-linux-gnu/
RUN mkdir /usr/local/nv_tools && tar -C /usr/local/nv_tools/ -xvf Linux_for_Tegra/nv_tegra/nv_tools.tbz2 home/nvidia --strip-components 2

# Setup paths
ENV PATH=$PATH:/usr/local/cuda/bin
ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/aarch64-linux-gnu/tegra

# Symlink libcuda.so.1 to the correct version of CUDA
RUN ln -s /usr/lib/aarch64-linux-gnu/tegra/libcuda.so.1.1 /usr/lib/aarch64-linux-gnu/tegra/libcuda.so.1

# Clean
RUN rm -rf *.deb *.tbz2 Linux_for_Tegra