Add pytorch-benchmark image

.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

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/

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