diff --git a/.github/workflows/publish-benchmark-images.yaml b/.github/workflows/publish-benchmark-images.yaml index a5a1bcf..20fbede 100644 --- a/.github/workflows/publish-benchmark-images.yaml +++ b/.github/workflows/publish-benchmark-images.yaml @@ -21,6 +21,7 @@ jobs: - component: mpi-benchmarks - component: openfoam - component: perftest + - component: pytorch-benchmarks steps: - name: Check out the repository uses: actions/checkout@v2 diff --git a/images/pytorch-benchmarks/Dockerfile b/images/pytorch-benchmarks/Dockerfile new file mode 100644 index 0000000..72fe8c4 --- /dev/null +++ b/images/pytorch-benchmarks/Dockerfile @@ -0,0 +1,13 @@ +# TODO: Try newer versions of base image - 23.08 fails with +# 'CUDA on host is too old' type error but something in between +# should work. Latest 23.08 might also work on Ubuntu-22.04. +FROM nvcr.io/nvidia/pytorch:22.12-py3 + +RUN pip install -U pip +RUN pip install torchvision torchaudio +RUN git clone https://github.com/pytorch/benchmark +WORKDIR /workspace/benchmark +# Add other benchmarks here? +RUN python install.py alexnet resnet50 llama + +COPY run-benchmark.sh /usr/local/bin/ \ No newline at end of file diff --git a/images/pytorch-benchmarks/run-benchmark.sh b/images/pytorch-benchmarks/run-benchmark.sh new file mode 100644 index 0000000..11cfcaa --- /dev/null +++ b/images/pytorch-benchmarks/run-benchmark.sh @@ -0,0 +1,3 @@ +#!/bin/bash +# Env vars should be set in container spec +python run.py $BENCHMARK_NAME -t $BENCHMARK_TYPE -d $BENCHMARK_DEVICE --bs $INPUT_DATA_BATCH_SIZE \ No newline at end of file