Dockerfile.pytorch

FROM timongentzsch/l4t-ubuntu20-opencv

ARG DEBIAN_FRONTEND=noninteractive
ARG BUILD_SOX=1
ARG PYTORCH_VERSION=1.11.0

# 
# install pytorch deps
#
RUN apt-get update && apt-get install -y --no-install-recommends python3-pip python3-dev libopenblas-dev libjpeg-dev zlib1g-dev sox libsox-dev libsox-fmt-all unzip && \
    rm -rf /var/lib/apt/lists/* && \
    apt-get clean

SHELL ["/bin/bash", "-c", "-l"]

# 
# install prebuild pytorch binaries
#
WORKDIR /tmp/
RUN --mount=target=/assets,type=bind,source=assets \
    cp /assets/torch-${PYTORCH_VERSION}-aarch64-ubuntu20.zip . && \
    unzip torch-${PYTORCH_VERSION}-aarch64-ubuntu20.zip && \
    pip3 install *.whl && \
    rm -rf *

# 
# install tensorrt python bindings
#
RUN --mount=target=/assets,type=bind,source=assets \
    cp /assets/tensorrt-8.2.1.8-cp38-none-linux_aarch64.whl . && \
    pip3 install *.whl && \
    rm -rf *

ENV PATH=/usr/local/cuda/bin:/usr/local/cuda-10.2/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda-10.2/targets/aarch64-linux/lib:

CMD ["bash"]
WORKDIR /root