Dockerfile for building Tensorflow 2.0.1 from source with Extra CPU (SSE4.1 SSE4.2 AVX AVX2 FMA) instructions

Dockerfile

FROM ubuntu:20.04

RUN apt update && apt install -y \
    build-essential \
    curl \
    git \
    wget \
    libc-ares-dev \
    libjpeg-dev \
    openjdk-8-jdk \
    gcc \
    g++ \
    python \
    python-pip \
    && rm -rf /var/lib/lists/*

ENV PYTHON_BIN_PATH="/usr/bin/python" \
    PYTHON_LIB_PATH="/usr/local/lib/python3.8/dist-packages" \
    PYTHONPATH=/tensorflow/lib \
    PYTHON_ARG=/tensorflow/lib \
    USE_BAZEL_VERSION=0.26.1 \
    TF_NEED_CUDA=0 \
    TF_NEED_GCP=0 \
    TF_CUDA_COMPUTE_CAPABILITIES=5.2,3.5 \
    TF_NEED_HDFS=0 \
    TF_NEED_OPENCL=0 \
    TF_NEED_JEMALLOC=0 \
    TF_ENABLE_XLA=0 \
    TF_NEED_VERBS=0 \
    TF_CUDA_CLANG=0 \
    TF_DOWNLOAD_CLANG=0 \
    TF_NEED_MKL=0 \
    TF_DOWNLOAD_MKL=0 \
    TF_NEED_MPI=0 \
    TF_NEED_S3=1 \
    TF_NEED_KAFKA=0 \
    TF_NEED_GDR=0 \
    TF_NEED_OPENCL_SYCL=0 \
    TF_SET_ANDROID_WORKSPACE=0 \
    TF_NEED_AWS=0 \
    TF_NEED_IGNITE=0 \
    TF_NEED_ROCM=0 \
    TF_VERSION=2.0.1\
    GCC_HOST_COMPILER_PATH="/usr/bin/gcc" \
    CC_OPT_FLAGS="-march=native"

RUN pip install pip six numpy wheel setuptools mock 'future>=0.17.1' && \
    pip install keras_applications --no-deps && \
    pip install keras_preprocessing --no-deps

RUN echo $TF_VERSION

RUN wget https://github.com/tensorflow/tensorflow/archive/v$TF_VERSION.tar.gz -q && \
    tar -xzf v$TF_VERSION.tar.gz && \
    mkdir /tensorflow && mv tensorflow-$TF_VERSION/* /tensorflow && \
    rm v$TF_VERSION.tar.gz

RUN wget https://github.com/bazelbuild/bazelisk/releases/download/v1.2.1/bazelisk-linux-amd64 -q && \
    mv bazelisk-linux-amd64 /bin/bazel && \
    chmod +x /bin/bazel

WORKDIR /tensorflow

RUN ./configure

RUN bazel build -c opt \
            --define=grpc_no_ares=true  \
            --linkopt="-lrt" \
            --linkopt="-lm" \
            --host_linkopt="-lrt" \
            --host_linkopt="-lm" \
            --action_env="LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" \
            --copt=-mavx --copt=-mavx2 --copt=-mfma --copt=-mfpmath=both \
            --copt=-w \
            --jobs=26 \
            //tensorflow/tools/lib_package:libtensorflow

Docker file for compiling TF2 with all CPU instructions available