Update pytorch benchmark Dockerfile.

Update pytorch benchmark to reflect a diversity of workloads.

This paper (https://arxiv.org/pdf/2304.14226.pdf) is written by the authors
of the used pytorch benchmarks. It details 1) what each benchmark's domain
is and the type of workload and 2) the profiles of each benchmark (e.g.
how much it taxes the GPU and if it moves data around).

In addition, the above benchmarks all have a significant runtime so that
we get more meaningful results with respect to overall performance WRT
both the GPU usage and the python applications themselves.

PiperOrigin-RevId: 547948747

images/benchmarks/pytorch/Dockerfile.x86_64

@@ -1,23 +1,34 @@
 FROM pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel
 
-RUN apt-get update
-RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git wget \
+RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y git wget \
     libgl1-mesa-glx libglib2.0-0
 
 RUN git clone https://github.com/pytorch/benchmark.git \
     && cd /workspace/benchmark \
     && git reset --hard 6d84522ec83997beefa4bb6e5154453b5fc653f7
 
+WORKDIR /workspace/benchmark
+
 # For some reason, installing the suite's requirements.txt does not correctly
 # install "boto3", so install it explicitly.
-RUN pip install boto3 && pip install -r /workspace/benchmark/requirements.txt
+# Note that mobilenet_v2 does not have a requirements.txt file.
+RUN pip install boto3 numba matplotlib && pip install \
+ -r requirements.txt \
+ -r torchbenchmark/models/LearningToPaint/requirements.txt \
+ -r torchbenchmark/models/attention_is_all_you_need_pytorch/requirements.txt \
+ -r torchbenchmark/models/fastNLP_Bert/requirements.txt \
+ -r torchbenchmark/models/hf_BigBird/requirements.txt \
+ -r torchbenchmark/models/speech_transformer/requirements.txt \
+ -r torchbenchmark/models/Background_Matting/requirements.txt
 
-RUN cd /workspace/benchmark && python install.py \
-    DALLE2_pytorch \
+# These benchmarks are chosen based on diversity of the type of model and their
+# profile with respect to using the GPU and moving data. For more context, see
+# this paper: https://arxiv.org/pdf/2304.14226.pdf
+RUN python install.py \
     LearningToPaint \
-    Super_SloMo
-
-RUN cd /workspace/benchmark/torchbenchmark/models && pip install \
-    -r DALLE2_pytorch/requirements.txt \
-    -r LearningToPaint/requirements.txt \
-    -r Super_SloMo/requirements.txt
+    attention_is_all_you_need_pytorch \
+    fastNLP_Bert \
+    hf_BigBird \
+    speech_transformer \
+    mobilenet_v2 \
+    Background_Matting