能否用Dockerfile编译支持SSE4.1、SSE4.2及GPU的TensorFlow二进制包？

Absolutely, you can build a TensorFlow binary that supports both SSE4.1/SSE4.2 instruction sets and GPU acceleration via a custom Dockerfile—this is totally feasible, and I’ll break down how to do it clearly:

1. Can a Dockerfile compile TensorFlow with both SSE4.1/SSE4.2 and GPU support? #

Yes, definitely. The core idea is combining two sets of configurations in the build process:

• SSE4.1/SSE4.2 support: Enabled via compiler flags (-msse4.1, -msse4.2) passed to TensorFlow’s build tool (Bazel) to optimize the CPU code path.
• GPU support: Enabled by linking against NVIDIA’s CUDA and cuDNN libraries, which we can get by starting with an official NVIDIA CUDA base image that includes development tools.

As long as you align compatible versions of CUDA, Bazel, and TensorFlow source code, you can enable both features in a single build—no need to choose one over the other. This directly solves the "works on host but fails in Docker" problem by ensuring the compiled binary matches your target environment’s hardware capabilities.

2. Dockerfile Example & Critical Setup Steps #

Here’s a production-ready Dockerfile that compiles TensorFlow with both features. Adjust version numbers to match your hardware and needs:

# Base image: NVIDIA CUDA + cuDNN + Ubuntu dev tools
ARG CUDA_VERSION=11.8.0
ARG CUDNN_VERSION=8
ARG UBUNTU_VERSION=22.04
FROM nvidia/cuda:${CUDA_VERSION}-cudnn${CUDNN_VERSION}-devel-ubuntu${UBUNTU_VERSION}

# Environment setup to avoid interactive prompts
ENV DEBIAN_FRONTEND=noninteractive
ENV TF_VERSION=r2.15  # Match to your desired TensorFlow version

# Install system dependencies for building
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    git \
    python3-dev \
    python3-pip \
    openjdk-11-jdk \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Install Bazel (must match TensorFlow's supported version; check TF docs)
RUN curl -fsSL https://github.com/bazelbuild/bazel/releases/download/6.4.0/bazel-6.4.0-installer-linux-x86_64.sh -o bazel-installer.sh \
    && chmod +x bazel-installer.sh \
    && ./bazel-installer.sh \
    && rm bazel-installer.sh

# Upgrade pip and install Python build dependencies
RUN pip3 install --upgrade pip setuptools wheel numpy==1.24.3

# Clone TensorFlow source code
RUN git clone --branch ${TF_VERSION} https://github.com/tensorflow/tensorflow.git /tensorflow
WORKDIR /tensorflow

# Configure TensorFlow build: enable CUDA + SSE4.x optimizations
RUN ./configure <<EOF
/usr/bin/python3
/tensorflow
NO
NO
YES
NO
NO
NO
--copt=-msse4.1 --copt=-msse4.2 --config=cuda
EOF

# Build TensorFlow (adjust --jobs based on your host's CPU cores to speed up compilation)
RUN bazel build --config=cuda --copt=-msse4.1 --copt=-msse4.2 //tensorflow/tools/pip_package:build_pip_package

# Create a pip package from the build
RUN ./bazel-bin/tensorflow/tools/pip_package/build_pip_package /tmp/tensorflow_pkg

# Install the compiled TensorFlow package
RUN pip3 install /tmp/tensorflow_pkg/tensorflow-*.whl

# Clean up build artifacts to reduce image size
RUN bazel clean --expunge && rm -rf /tensorflow /tmp/tensorflow_pkg

Key Tips for Success: #

• Base Image Choice: Use the devel variant of the NVIDIA CUDA image—it includes the headers and libraries needed to compile GPU-supporting code (the runtime-only variant won’t work here).
• Version Alignment: Double-check that your Bazel version matches the one required by your target TensorFlow version (TensorFlow’s official build docs list compatible versions).
• Build Resources: Compiling TensorFlow is CPU/GPU intensive—expect this to take 30+ minutes. Use the --jobs flag in the bazel build command to match your host’s core count.
• Host Compatibility: Your host machine needs the NVIDIA Docker Runtime installed, and its CUDA version should be compatible with the one in the Docker image.

Validate the Build #

After building the image, run this command to confirm both features work:

docker run --gpus all your-tf-image python3 -c "
import tensorflow as tf
# Check GPU availability
print('GPU Available:', tf.test.is_gpu_available())
# Verify SSE4.x support
from tensorflow.python.framework import test_util
print('SSE4.1 Supported:', test_util.IsSSE41Supported())
print('SSE4.2 Supported:', test_util.IsSSE42Supported())
"

内容的提问来源于stack exchange，提问作者Diego Orellana