Build tflite_runtime Python wheel with CMake

[Normal build for Python 3.7] $ PYTHON=python3.7 tensorflow/lite/tools/pip_package/build_pip_package_with_cmake.sh [armhf cross build for Python 3.7] $ tensorflow/tools/ci_build/ci_build.sh PI-PYTHON37 tensorflow/lite/tools/pip_package/build_pip_package_with_cmake.sh armhf - lite/tools/cmake/download_toolchains.sh is added to download ARM toolchains - tools/ci_build/Dockerfile.pi* are updated to use cmake 3.16.8 PiperOrigin-RevId: 356621208 Change-Id: I7f1cb0c8389b5a5e637036b52b33e5da0afa2b77
2021-02-09 16:37:08 -08:00 · 2021-02-09 16:37:08 -08:00 · e54a4e3aea
commit e54a4e3aea
parent 6b8040d39c
10 changed files with 346 additions and 16 deletions
--- a/tensorflow/lite/CMakeLists.txt
+++ b/tensorflow/lite/CMakeLists.txt
@ -38,6 +38,7 @@ cmake_policy(SET CMP0028 NEW)
 # Enable MACOSX_RPATH (@rpath) for built dynamic libraries.
 cmake_policy(SET CMP0042 NEW)
 project(tensorflow-lite C CXX)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 set(TENSORFLOW_SOURCE_DIR "" CACHE PATH
  "Directory that contains the TensorFlow project"
 )
@ -488,3 +489,28 @@ target_link_libraries(benchmark_model
    ${TFLITE_BENCHMARK_LIBS}
 )

+# Python interpreter wrapper.
+add_library(_pywrap_tensorflow_interpreter_wrapper SHARED EXCLUDE_FROM_ALL
+  ${TFLITE_SOURCE_DIR}/python/interpreter_wrapper/interpreter_wrapper.cc
+  ${TFLITE_SOURCE_DIR}/python/interpreter_wrapper/interpreter_wrapper_pybind11.cc
+  ${TFLITE_SOURCE_DIR}/python/interpreter_wrapper/numpy.cc
+  ${TFLITE_SOURCE_DIR}/python/interpreter_wrapper/python_error_reporter.cc
+  ${TFLITE_SOURCE_DIR}/python/interpreter_wrapper/python_utils.cc
+)
+
+# To remove "lib" prefix.
+set_target_properties(_pywrap_tensorflow_interpreter_wrapper PROPERTIES PREFIX "")
+
+target_include_directories(_pywrap_tensorflow_interpreter_wrapper
+  PUBLIC
+    ${TFLITE_INCLUDE_DIRS}
+)
+
+target_link_libraries(_pywrap_tensorflow_interpreter_wrapper
+  tensorflow-lite
+  ${CMAKE_DL_LIBS}
+)
+target_compile_options(_pywrap_tensorflow_interpreter_wrapper
+  PUBLIC ${TFLITE_TARGET_PUBLIC_OPTIONS}
+  PRIVATE ${TFLITE_TARGET_PRIVATE_OPTIONS}
+)
--- a/tensorflow/lite/tools/cmake/download_toolchains.sh
+++ b/tensorflow/lite/tools/cmake/download_toolchains.sh
@ -0,0 +1,78 @@
+#!/bin/bash
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+cd "$SCRIPT_DIR/../../../.."
+
+TOOLCHAINS_DIR=$(realpath tensorflow/lite/tools/cmake/toolchains)
+mkdir -p ${TOOLCHAINS_DIR}
+
+case $1 in
+	armhf)
+    if [[ ! -d "${TOOLCHAINS_DIR}/gcc-arm-8.3-2019.03-x86_64-arm-linux-gnueabihf" ]]; then
+      curl -LO https://storage.googleapis.com/mirror.tensorflow.org/developer.arm.com/media/Files/downloads/gnu-a/8.3-2019.03/binrel/gcc-arm-8.3-2019.03-x86_64-arm-linux-gnueabihf.tar.xz >&2
+      tar xvf gcc-arm-8.3-2019.03-x86_64-arm-linux-gnueabihf.tar.xz -C ${TOOLCHAINS_DIR} >&2
+    fi
+    ARMCC_ROOT=${TOOLCHAINS_DIR}/gcc-arm-8.3-2019.03-x86_64-arm-linux-gnueabihf
+    echo "ARMCC_FLAGS=\"-march=armv7-a -mfpu=neon-vfpv4 -funsafe-math-optimizations \
+      -isystem ${ARMCC_ROOT}/lib/gcc/arm-linux-gnueabihf/8.3.0/include \
+      -isystem ${ARMCC_ROOT}/lib/gcc/arm-linux-gnueabihf/8.3.0/include-fixed \
+      -isystem ${ARMCC_ROOT}/arm-linux-gnueabihf/include/c++/8.3.0 \
+      -isystem ${ARMCC_ROOT}/arm-linux-gnueabihf/libc/usr/include \
+      -isystem \"\${CROSSTOOL_PYTHON_INCLUDE_PATH}\" \
+      -isystem /usr/include\""
+    echo "ARMCC_PREFIX=${ARMCC_ROOT}/bin/arm-linux-gnueabihf-"
+		;;
+	aarch64)
+    if [[ ! -d "${TOOLCHAINS_DIR}/gcc-arm-8.3-2019.03-x86_64-aarch64-linux-gnu" ]]; then
+      curl -LO https://storage.googleapis.com/mirror.tensorflow.org/developer.arm.com/media/Files/downloads/gnu-a/8.3-2019.03/binrel/gcc-arm-8.3-2019.03-x86_64-aarch64-linux-gnu.tar.xz >&2
+      tar xvf gcc-arm-8.3-2019.03-x86_64-aarch64-linux-gnu.tar.xz -C ${TOOLCHAINS_DIR} >&2
+    fi
+    ARMCC_ROOT=${TOOLCHAINS_DIR}/gcc-arm-8.3-2019.03-x86_64-aarch64-linux-gnu
+    echo "ARMCC_FLAGS=\"-funsafe-math-optimizations \
+      -isystem ${ARMCC_ROOT}/lib/gcc/aarch64-linux-gnu/8.3.0/include \
+      -isystem ${ARMCC_ROOT}/lib/gcc/aaarch64-linux-gnu/8.3.0/include-fixed \
+      -isystem ${ARMCC_ROOT}/aarch64-linux-gnu/include/c++/8.3.0 \
+      -isystem ${ARMCC_ROOT}/aarch64-linux-gnu/libc/usr/include \
+      -isystem \"\${CROSSTOOL_PYTHON_INCLUDE_PATH}\" \
+      -isystem /usr/include\""
+    echo "ARMCC_PREFIX=${ARMCC_ROOT}/bin/aarch64-linux-gnu-"
+		;;
+	rpi0)
+    if [[ ! -d "${TOOLCHAINS_DIR}/arm-rpi-linux-gnueabihf" ]]; then
+      curl -L https://github.com/rvagg/rpi-newer-crosstools/archive/eb68350c5c8ec1663b7fe52c742ac4271e3217c5.tar.gz -o rpi-toolchain.tar.gz >&2
+      tar xzf rpi-toolchain.tar.gz -C ${TOOLCHAINS_DIR} >&2
+      mv ${TOOLCHAINS_DIR}/rpi-newer-crosstools-eb68350c5c8ec1663b7fe52c742ac4271e3217c5 ${TOOLCHAINS_DIR}/arm-rpi-linux-gnueabihf >&2
+    fi
+    ARMCC_ROOT=${TOOLCHAINS_DIR}/arm-rpi-linux-gnueabihf/x64-gcc-6.5.0/arm-rpi-linux-gnueabihf
+    echo "ARMCC_FLAGS=\"-march=armv6 -mfpu=vfp -funsafe-math-optimizations \
+      -isystem ${ARMCC_ROOT}/lib/gcc/arm-rpi-gnueabihf/6.5.0/include \
+      -isystem ${ARMCC_ROOT}/lib/gcc/arm-rpi-gnueabihf/6.5.0/include-fixed \
+      -isystem ${ARMCC_ROOT}/arm-rpi-linux-gnueabihf/include/c++/6.5.0 \
+      -isystem ${ARMCC_ROOT}/arm-rpi-linux-gnueabihf/sysroot/usr/include \
+      -isystem \"\${CROSSTOOL_PYTHON_INCLUDE_PATH}\" \
+      -isystem /usr/include\""
+    echo "ARMCC_PREFIX=${ARMCC_ROOT}/bin/arm-rpi-linux-gnueabihf-"
+    ;;
+	*)
+		echo "Usage: download_toolchains.sh [armhf|aarch64|rpi0]" >&2
+    exit
+		;;
+  esac
+
+echo "download_toolchains.sh completed successfully." >&2
--- a/tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh
+++ b/tensorflow/lite/tools/pip_package/build_pip_package_with_bazel.sh
@ -43,7 +43,7 @@ fi
 # Build source tree.
 rm -rf "${BUILD_DIR}" && mkdir -p "${BUILD_DIR}/tflite_runtime"
 cp -r "${TENSORFLOW_LITE_DIR}/tools/pip_package/debian" \
-      "${TENSORFLOW_LITE_DIR}/tools/pip_package/setup_with_bazel.py" \
+      "${TENSORFLOW_LITE_DIR}/tools/pip_package/setup_with_binary.py" \
      "${TENSORFLOW_LITE_DIR}/tools/pip_package/MANIFEST.in" \
      "${TENSORFLOW_LITE_DIR}/python/interpreter_wrapper" \
      "${BUILD_DIR}"
@ -101,19 +101,19 @@ chmod u+w "${BUILD_DIR}/tflite_runtime/_pywrap_tensorflow_interpreter_wrapper${L
 cd "${BUILD_DIR}"
 case "${TENSORFLOW_TARGET}" in
  armhf)
-    ${PYTHON} setup_with_bazel.py bdist --plat-name=linux-armv7l \
+    ${PYTHON} setup_with_binary.py bdist --plat-name=linux-armv7l \
                       bdist_wheel --plat-name=linux-armv7l
    ;;
  aarch64)
-    ${PYTHON} setup_with_bazel.py bdist --plat-name=linux-aarch64 \
+    ${PYTHON} setup_with_binary.py bdist --plat-name=linux-aarch64 \
                       bdist_wheel --plat-name=linux-aarch64
    ;;
  *)
    if [[ -n "${TENSORFLOW_TARGET}" ]] && [[ -n "${TENSORFLOW_TARGET_ARCH}" ]]; then
-      ${PYTHON} setup_with_bazel.py bdist --plat-name=${TENSORFLOW_TARGET}-${TENSORFLOW_TARGET_ARCH} \
+      ${PYTHON} setup_with_binary.py bdist --plat-name=${TENSORFLOW_TARGET}-${TENSORFLOW_TARGET_ARCH} \
                         bdist_wheel --plat-name=${TENSORFLOW_TARGET}-${TENSORFLOW_TARGET_ARCH}
    else
-      ${PYTHON} setup_with_bazel.py bdist bdist_wheel
+      ${PYTHON} setup_with_binary.py bdist bdist_wheel
    fi
    ;;
 esac
--- a/tensorflow/lite/tools/pip_package/build_pip_package_with_cmake.sh
+++ b/tensorflow/lite/tools/pip_package/build_pip_package_with_cmake.sh
@ -0,0 +1,203 @@
+#!/usr/bin/env bash
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+set -ex
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PYTHON="${CI_BUILD_PYTHON:-python3}"
+VERSION_SUFFIX=${VERSION_SUFFIX:-}
+export TENSORFLOW_DIR="${SCRIPT_DIR}/../../../.."
+TENSORFLOW_LITE_DIR="${TENSORFLOW_DIR}/tensorflow/lite"
+TENSORFLOW_VERSION=$(grep "_VERSION = " "${TENSORFLOW_DIR}/tensorflow/tools/pip_package/setup.py" | cut -d= -f2 | sed "s/[ '-]//g")
+export PACKAGE_VERSION="${TENSORFLOW_VERSION}${VERSION_SUFFIX}"
+BUILD_DIR="${SCRIPT_DIR}/gen/tflite_pip/${PYTHON}"
+TENSORFLOW_TARGET=$1
+PYTHON_INCLUDE=$(${PYTHON} -c "from sysconfig import get_paths as gp; print(gp()['include'])")
+PYBIND11_INCLUDE=$(${PYTHON} -c "import pybind11; print (pybind11.get_include())")
+
+# Fix container image for cross build.
+if [ ! -z "${CI_BUILD_HOME}" ] && [ `pwd` = "/workspace" ]; then
+  # Fix for curl build problem in 32-bit, see https://stackoverflow.com/questions/35181744/size-of-array-curl-rule-01-is-negative
+  if [ "${TENSORFLOW_TARGET}" = "armhf" ]; then
+    sudo sed -i 's/define CURL_SIZEOF_LONG 8/define CURL_SIZEOF_LONG 4/g' /usr/include/curl/curlbuild.h
+    sudo sed -i 's/define CURL_SIZEOF_CURL_OFF_T 8/define CURL_SIZEOF_CURL_OFF_T 4/g' /usr/include/curl/curlbuild.h
+  fi
+
+  # The system-installed OpenSSL headers get pulled in by the latest BoringSSL
+  # release on this configuration, so move them before we build:
+  if [ -d /usr/include/openssl ]; then
+    sudo mv /usr/include/openssl /usr/include/openssl.original
+  fi
+fi
+
+# Build source tree.
+rm -rf "${BUILD_DIR}" && mkdir -p "${BUILD_DIR}/tflite_runtime"
+cp -r "${TENSORFLOW_LITE_DIR}/tools/pip_package/debian" \
+      "${TENSORFLOW_LITE_DIR}/tools/pip_package/setup_with_binary.py" \
+      "${TENSORFLOW_LITE_DIR}/tools/pip_package/MANIFEST.in" \
+      "${TENSORFLOW_LITE_DIR}/python/interpreter_wrapper" \
+      "${BUILD_DIR}"
+cp "${TENSORFLOW_LITE_DIR}/python/interpreter.py" \
+   "${BUILD_DIR}/tflite_runtime"
+echo "__version__ = '${PACKAGE_VERSION}'" >> "${BUILD_DIR}/tflite_runtime/__init__.py"
+echo "__git_version__ = '$(git -C "${TENSORFLOW_DIR}" describe)'" >> "${BUILD_DIR}/tflite_runtime/__init__.py"
+
+# Build python interpreter_wrapper.
+mkdir -p "${BUILD_DIR}/cmake_build"
+cd "${BUILD_DIR}/cmake_build"
+
+echo "Building for ${TENSORFLOW_TARGET}"
+case "${TENSORFLOW_TARGET}" in
+  armhf)
+    eval $(${TENSORFLOW_LITE_DIR}/tools/cmake/download_toolchains.sh "${TENSORFLOW_TARGET}")
+    ARMCC_FLAGS="${ARMCC_FLAGS} -I${PYBIND11_INCLUDE}"
+    cmake \
+      -DCMAKE_C_COMPILER=${ARMCC_PREFIX}gcc \
+      -DCMAKE_CXX_COMPILER=${ARMCC_PREFIX}g++ \
+      -DCMAKE_C_FLAGS="${ARMCC_FLAGS}" \
+      -DCMAKE_CXX_FLAGS="${ARMCC_FLAGS}" \
+      -DCMAKE_SYSTEM_NAME=Linux \
+      -DCMAKE_SYSTEM_PROCESSOR=armv7 \
+      "${TENSORFLOW_LITE_DIR}"
+    ;;
+  rpi0)
+    eval $(${TENSORFLOW_LITE_DIR}/tools/cmake/download_toolchains.sh "${TENSORFLOW_TARGET}")
+    ARMCC_FLAGS="${ARMCC_FLAGS} -I${PYBIND11_INCLUDE}"
+    cmake \
+      -DCMAKE_C_COMPILER=${ARMCC_PREFIX}gcc \
+      -DCMAKE_CXX_COMPILER=${ARMCC_PREFIX}g++ \
+      -DCMAKE_C_FLAGS="${ARMCC_FLAGS}" \
+      -DCMAKE_CXX_FLAGS="${ARMCC_FLAGS}" \
+      -DCMAKE_SYSTEM_NAME=Linux \
+      -DCMAKE_SYSTEM_PROCESSOR=armv6 \
+      -DTFLITE_ENABLE_XNNPACK=OFF \
+      "${TENSORFLOW_LITE_DIR}"
+    ;;
+  aarch64)
+    eval $(${TENSORFLOW_LITE_DIR}/tools/cmake/download_toolchains.sh "${TENSORFLOW_TARGET}")
+    ARMCC_FLAGS="${ARMCC_FLAGS} -I${PYBIND11_INCLUDE}"
+    cmake \
+      -DCMAKE_C_COMPILER=${ARMCC_PREFIX}gcc \
+      -DCMAKE_CXX_COMPILER=${ARMCC_PREFIX}g++ \
+      -DCMAKE_C_FLAGS="${ARMCC_FLAGS}" \
+      -DCMAKE_CXX_FLAGS="${ARMCC_FLAGS}" \
+      -DCMAKE_SYSTEM_NAME=Linux \
+      -DCMAKE_SYSTEM_PROCESSOR=aarch64 \
+      "${TENSORFLOW_LITE_DIR}"
+    ;;
+  native)
+    BUILD_FLAGS=${BUILD_FLAGS:-"-march=native -I${PYTHON_INCLUDE} -I${PYBIND11_INCLUDE}"}
+    cmake \
+      -DCMAKE_C_FLAGS="${BUILD_FLAGS}" \
+      -DCMAKE_CXX_FLAGS="${BUILD_FLAGS}" \
+      "${TENSORFLOW_LITE_DIR}"
+    ;;
+  *)
+    BUILD_FLAGS=${BUILD_FLAGS:-"-I${PYTHON_INCLUDE} -I${PYBIND11_INCLUDE}"}
+    cmake \
+      -DCMAKE_C_FLAGS="${BUILD_FLAGS}" \
+      -DCMAKE_CXX_FLAGS="${BUILD_FLAGS}" \
+      "${TENSORFLOW_LITE_DIR}"
+    ;;
+esac
+
+cmake --build . --verbose -j -t _pywrap_tensorflow_interpreter_wrapper
+cd "${BUILD_DIR}"
+
+case "${TENSORFLOW_TARGET}" in
+  windows)
+    LIBRARY_EXTENSION=".pyd"
+    ;;
+  *)
+    LIBRARY_EXTENSION=".so"
+    ;;
+esac
+
+cp "${BUILD_DIR}/cmake_build/_pywrap_tensorflow_interpreter_wrapper${LIBRARY_EXTENSION}" \
+   "${BUILD_DIR}/tflite_runtime"
+# Bazel generates the wrapper library with r-x permissions for user.
+# At least on Windows, we need write permissions to delete the file.
+# Without this, setuptools fails to clean the build directory.
+chmod u+w "${BUILD_DIR}/tflite_runtime/_pywrap_tensorflow_interpreter_wrapper${LIBRARY_EXTENSION}"
+
+# Build python wheel.
+cd "${BUILD_DIR}"
+case "${TENSORFLOW_TARGET}" in
+  armhf)
+    ${PYTHON} setup_with_binary.py bdist --plat-name=linux-armv7l \
+                       bdist_wheel --plat-name=linux-armv7l
+    ;;
+  rpi0)
+    ${PYTHON} setup_with_binary.py bdist --plat-name=linux_armv6l \
+                       bdist_wheel --plat-name=linux-armv6l
+    ;;
+  aarch64)
+    ${PYTHON} setup_with_binary.py bdist --plat-name=linux-aarch64 \
+                       bdist_wheel --plat-name=linux-aarch64
+    ;;
+  *)
+    if [[ -n "${TENSORFLOW_TARGET}" ]] && [[ -n "${TENSORFLOW_TARGET_ARCH}" ]]; then
+      ${PYTHON} setup_with_binary.py bdist --plat-name=${TENSORFLOW_TARGET}-${TENSORFLOW_TARGET_ARCH} \
+                         bdist_wheel --plat-name=${TENSORFLOW_TARGET}-${TENSORFLOW_TARGET_ARCH}
+    else
+      ${PYTHON} setup_with_binary.py bdist bdist_wheel
+    fi
+    ;;
+esac
+
+echo "Output can be found here:"
+find "${BUILD_DIR}"
+
+# Build debian package.
+if [[ "${BUILD_DEB}" != "y" ]]; then
+  exit 0
+fi
+
+PYTHON_VERSION=$(${PYTHON} -c "import sys;print(sys.version_info.major)")
+if [[ ${PYTHON_VERSION} != 3 ]]; then
+  echo "Debian package can only be generated for python3." >&2
+  exit 1
+fi
+
+DEB_VERSION=$(dpkg-parsechangelog --show-field Version | cut -d- -f1)
+if [[ "${DEB_VERSION}" != "${PACKAGE_VERSION}" ]]; then
+  cat << EOF > "${BUILD_DIR}/debian/changelog"
+tflite-runtime (${PACKAGE_VERSION}-1) unstable; urgency=low
+
+  * Bump version to ${PACKAGE_VERSION}.
+
+ -- TensorFlow team <packages@tensorflow.org>  $(date -R)
+
+$(<"${BUILD_DIR}/debian/changelog")
+EOF
+fi
+
+case "${TENSORFLOW_TARGET}" in
+  armhf)
+    dpkg-buildpackage -b -rfakeroot -us -uc -tc -d -a armhf
+    ;;
+  rpi0)
+    dpkg-buildpackage -b -rfakeroot -us -uc -tc -d -a armel
+    ;;
+  aarch64)
+    dpkg-buildpackage -b -rfakeroot -us -uc -tc -d -a arm64
+    ;;
+  *)
+    dpkg-buildpackage -b -rfakeroot -us -uc -tc -d
+    ;;
+esac
+
+cat "${BUILD_DIR}/debian/changelog"
+
--- a/tensorflow/lite/tools/pip_package/setup_with_binary.py
+++ b/tensorflow/lite/tools/pip_package/setup_with_binary.py
@ -67,5 +67,5 @@ setup(
    package_dir={'': '.'},
    package_data={'': ['*.so', '*.pyd']},
    install_requires=[
-        'numpy >= 1.16.0',
+        'numpy ~= 1.19.2',  # Higher versions have a compatibility issue.
    ])
--- a/tensorflow/tools/ci_build/Dockerfile.pi-python3
+++ b/tensorflow/tools/ci_build/Dockerfile.pi-python3
@ -8,9 +8,9 @@ ENV CROSSTOOL_PYTHON_INCLUDE_PATH=/usr/include/python3.5
 # Copy and run the install scripts.
 COPY install/*.sh /install/
 RUN /install/install_bootstrap_deb_packages.sh
-RUN add-apt-repository -y ppa:openjdk-r/ppa && \
-    add-apt-repository -y ppa:george-edison55/cmake-3.x
-RUN /install/install_deb_packages.sh
+RUN add-apt-repository -y ppa:openjdk-r/ppa
+RUN /install/install_deb_packages.sh --without_cmake
+RUN /install/install_cmake.sh
 RUN /install/install_pip_packages.sh
 RUN /install/install_bazel.sh
 RUN /install/install_proto3.sh
--- a/tensorflow/tools/ci_build/Dockerfile.pi-python37
+++ b/tensorflow/tools/ci_build/Dockerfile.pi-python37
@ -8,9 +8,9 @@ ENV CROSSTOOL_PYTHON_INCLUDE_PATH=/usr/include/python3.7
 # Copy and run the install scripts.
 COPY install/*.sh /install/
 RUN /install/install_bootstrap_deb_packages.sh
-RUN add-apt-repository -y ppa:openjdk-r/ppa && \
-    add-apt-repository -y ppa:george-edison55/cmake-3.x
-RUN /install/install_deb_packages.sh
+RUN add-apt-repository -y ppa:openjdk-r/ppa
+RUN /install/install_deb_packages.sh --without_cmake
+RUN /install/install_cmake.sh

 # The following line installs the Python 3.7 cross-compilation toolchain.
 RUN /install/install_pi_python3x_toolchain.sh "3.7"
--- a/tensorflow/tools/ci_build/Dockerfile.pi-python38
+++ b/tensorflow/tools/ci_build/Dockerfile.pi-python38
@ -8,9 +8,9 @@ ENV CROSSTOOL_PYTHON_INCLUDE_PATH=/usr/include/python3.8
 # Copy and run the install scripts.
 COPY install/*.sh /install/
 RUN /install/install_bootstrap_deb_packages.sh
-RUN add-apt-repository -y ppa:openjdk-r/ppa && \
-    add-apt-repository -y ppa:george-edison55/cmake-3.x
-RUN /install/install_deb_packages.sh
+RUN add-apt-repository -y ppa:openjdk-r/ppa
+RUN /install/install_deb_packages.sh --without_cmake
+RUN /install/install_cmake.sh

 # The following line installs the Python 3.8 cross-compilation toolchain.
 RUN /install/install_pi_python3x_toolchain.sh "3.8"
--- a/tensorflow/tools/ci_build/install/install_cmake.sh
+++ b/tensorflow/tools/ci_build/install/install_cmake.sh
@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+curl -OL https://github.com/Kitware/CMake/releases/download/v3.16.8/cmake-3.16.8-Linux-x86_64.sh
+echo "0241a05bee0dcdf60e912057cc86cbedba21b9b0d67ec11bc67ad4834f182a23 cmake-3.16.8-Linux-x86_64.sh" | sha256sum -c
+sh cmake-3.16.8-Linux-x86_64.sh --prefix=/usr --skip-license
--- a/tensorflow/tools/ci_build/install/install_pip_packages_by_version.sh
+++ b/tensorflow/tools/ci_build/install/install_pip_packages_by_version.sh
@ -26,6 +26,10 @@ if [[ ! -x "$(which "${PIP}")" ]]; then
 fi

 PACKAGES=(
+  # NOTE: As numpy has releases that break semver guarantees and several other
+  # deps depend on numpy without an upper bound, we must install numpy before
+  # everything else.
+  "numpy ~= 1.19.2"
  "auditwheel"
  "wheel"
  "setuptools"
@ -37,7 +41,6 @@ PACKAGES=(
  "bleach"
  "markdown"
  "protobuf"
-  "numpy"
  "scipy"
  "scikit-learn"
  "pandas"
@ -59,6 +62,7 @@ PACKAGES=(
  "dm-tree"
  "dill"
  "tblib"
+  "pybind11"
 )

 # tf.mock require the following for python2: