From 20ec1a3fa6a61a95bcac158b84f3035724c95b39 Mon Sep 17 00:00:00 2001 From: Dustin Herbison Date: Thu, 4 Jan 2024 15:44:41 +0000 Subject: [PATCH] Install Keras 3 http://b/318672157 --- tpu/Dockerfile | 5 +++++ tpu/config.txt | 6 +++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/tpu/Dockerfile b/tpu/Dockerfile index e072e403..fdc75de6 100644 --- a/tpu/Dockerfile +++ b/tpu/Dockerfile @@ -62,6 +62,11 @@ RUN pip install tensorflow_hub https://storage.googleapis.com/cloud-tpu-tpuvm-ar kagglehub && \ /tmp/clean-layer.sh +# b/318672157 Install Keras 3 for loading models. +# We ignore dependencies for now since tf2.15 should work despite pip saying it won't. +RUN pip install --no-dependencies "keras>3" keras-cv keras-nlp && \ + /tmp/clean-layer.sh + # Tensorflow libtpu: RUN curl --output /lib/libtpu.so https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/${TF_LIBTPU_VERSION}/libtpu.so diff --git a/tpu/config.txt b/tpu/config.txt index 86b097cc..a56a939d 100644 --- a/tpu/config.txt +++ b/tpu/config.txt @@ -1,9 +1,9 @@ BASE_IMAGE=python:3.10 PYTHON_WHEEL_VERSION=cp310 PYTHON_VERSION_PATH=python3.10 -# https://cloud.google.com/tpu/docs/supported-tpu-configurations#tpu_software_versions:~:text=TensorFlow%20version-,libtpu.so%20version,-2.13.0 -TENSORFLOW_VERSION=2.14.0 -TF_LIBTPU_VERSION=1.8.0 +# https://cloud.google.com/tpu/docs/supported-tpu-configurations#libtpu_versions +TENSORFLOW_VERSION=2.15.0 +TF_LIBTPU_VERSION=1.9.0 TF_LINUX_WHEEL_VERSION=manylinux_2_17_x86_64.manylinux2014_x86_64 JAX_VERSION=0.4.17 # Supports nightly