Update Dockerfile.pip

Author: EikanWang

docker/Dockerfile.pip

@@ -20,12 +20,15 @@ RUN apt-get update -y && \
     apt-get upgrade -y && \
     apt-get install -y --no-install-recommends --fix-missing \
     ${PYTHON} \
-    ${PYTHON}-pip
+    ${PYTHON}-pip \
+    google-perftools \
+    libjemalloc-dev \
+    numactl
 
 RUN ${PYTHON} -m pip --no-cache-dir install --upgrade \
     pip \
-    setuptools \
-    psutil
+    psutil \
+    setuptools
 
 # Some TF tools expect a "python" binary
 RUN ln -s $(which ${PYTHON}) /usr/local/bin/python
@@ -37,8 +40,12 @@ ARG TORCHVISION_VERSION=0.13.0+cpu
 ARG TORCH_CPU_URL=https://download.pytorch.org/whl/cpu/torch_stable.html
 ARG IPEX_URL=https://software.intel.com/ipex-whl-stable
 
-RUN \
-    python -m pip install --no-cache-dir \
+RUN python -m pip install --no-cache-dir \
+    intel-openmp \
     torch==${PYTORCH_VERSION} torchvision==${TORCHVISION_VERSION} torchaudio==${TORCHAUDIO_VERSION} -f ${TORCH_CPU_URL} && \
     python -m pip install --no-cache-dir \
     intel_extension_for_pytorch==${IPEX_VERSION} -f ${IPEX_URL}
+
+RUN ln -sf /usr/lib/x86_64-linux-gnu/libjemalloc.so /usr/lib/x86_64-linux-gnu/libtcmalloc.so
+ENV DNNL_MAX_CPU_ISA="AVX512_CORE_AMX"
+ENV LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libjemalloc.so:/usr/lib/x86_64-linux-gnu/libtcmalloc.so:/usr/local/lib/libiomp5.so":${LD_PRELOAD}