cjones/immich
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ml): introduce support of onnxruntime-rocm for AMD GPU

Browse Source
This commit is contained in:
Mehdi GHESH
2024-07-13 00:40:29 +02:00
committed by mertalev
parent 3f4bbab4eb
commit fe26ccd1b7
14 changed files with 275 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
machine-learning/Dockerfile
+42 -1
View File
@@ -15,6 +15,40 @@ RUN mkdir /opt/armnn && \
cd /opt/ann && \
sh build.sh
# Warning: 26.3Gb of disk space required to pull this image
# https://github.com/microsoft/onnxruntime/blob/main/dockerfiles/Dockerfile.rocm
FROM rocm/dev-ubuntu-22.04:6.1.2-complete as builder-rocm
WORKDIR /code
RUN apt-get update && apt-get install -y --no-install-recommends wget git python3.10-venv
# Install same version as the Dockerfile provided by onnxruntime
RUN wget https://github.com/Kitware/CMake/releases/download/v3.27.3/cmake-3.27.3-linux-x86_64.sh && \
chmod +x cmake-3.27.3-linux-x86_64.sh && \
mkdir -p /code/cmake-3.27.3-linux-x86_64 && \
./cmake-3.27.3-linux-x86_64.sh --skip-license --prefix=/code/cmake-3.27.3-linux-x86_64 && \
rm cmake-3.27.3-linux-x86_64.sh
ENV PATH /code/cmake-3.27.3-linux-x86_64/bin:${PATH}
# Prepare onnxruntime repository & build onnxruntime
RUN git clone --single-branch --branch v1.18.1 --recursive "https://github.com/Microsoft/onnxruntime" onnxruntime
WORKDIR /code/onnxruntime
# EDIT PR
# While there's still this PR open, we need to compile on the branch of the PR
# https://github.com/microsoft/onnxruntime/pull/19567
COPY ./rocm-PR19567.patch /tmp/
RUN git apply /tmp/rocm-PR19567.patch
# END EDIT PR
RUN /bin/sh ./dockerfiles/scripts/install_common_deps.sh
# I ran into a compilation error when parallelizing the build
# I used 12 threads to build onnxruntime, but it needs more than 16GB of RAM, and that's the amount of RAM I have on my machine
# I lowered the number of threads to 8, and it worked
# Even with 12 threads, the compilation took more than 1,5 hours to fail
RUN ./build.sh --allow_running_as_root --config Release --build_wheel --update --build --parallel 9 --cmake_extra_defines\
ONNXRUNTIME_VERSION=1.18.1 --use_rocm --rocm_home=/opt/rocm
RUN mv /code/onnxruntime/build/Linux/Release/dist/*.whl /opt/
FROM builder-${DEVICE} AS builder
ARG DEVICE
@@ -32,6 +66,9 @@ RUN poetry config installer.max-workers 10 && \
RUN python3 -m venv /opt/venv
COPY poetry.lock pyproject.toml ./
RUN if [ "$DEVICE" = "rocm" ]; then \
poetry add /opt/onnxruntime_rocm-*.whl; \
fi
RUN poetry install --sync --no-interaction --no-ansi --no-root --with ${DEVICE} --without dev
FROM python:3.11-slim-bookworm@sha256:614c8691ab74150465ec9123378cd4dde7a6e57be9e558c3108df40664667a4c AS prod-cpu
@@ -80,11 +117,15 @@ COPY --from=builder-armnn \
/opt/ann/build.sh \
/opt/armnn/
FROM rocm/dev-ubuntu-22.04:6.1.2-complete AS prod-rocm
FROM prod-${DEVICE} AS prod
ARG DEVICE
RUN apt-get update && \
apt-get install -y --no-install-recommends tini $(if ! [ "$DEVICE" = "openvino" ]; then echo "libmimalloc2.0"; fi) && \
apt-get install -y --no-install-recommends tini $(if ! [ "$DEVICE" = "openvino" ] && ! [ "$DEVICE" = "rocm" ]; then echo "libmimalloc2.0"; fi) && \
apt-get autoremove -yqq && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*