Replace tf.keras with tf_keras.

PiperOrigin-RevId: 601706370
tensorflow · Jan 26, 2024 · b9a4cf1 · b9a4cf1
1 parent 207b647
commit b9a4cf1
Show file tree

Hide file tree

Showing 33 changed files with 617 additions and 335 deletions.
diff --git a/configure/setup.py b/configure/setup.py
@@ -34,6 +34,7 @@
     "absl_py",
     "wheel",
     "wurlitzer",
+    "tf_keras",
 ]
 
 

diff --git a/documentation/known_issues.md b/documentation/known_issues.md
@@ -90,7 +90,7 @@ to the estimator format.
 
 While abstracted by the Keras API, a model instantiated in Python (e.g., with
 `tfdf.keras.RandomForestModel()`) and a model loaded from disk (e.g., with
-`tf.keras.models.load_model()`) can behave differently. Notably, a Python
+`tf_keras.models.load_model()`) can behave differently. Notably, a Python
 instantiated model automatically applies necessary type conversions. For
 example, if a `float64` feature is fed to a model expecting a `float32` feature,
 this conversion is performed implicitly. However, such a conversion is not

diff --git a/documentation/migration.md b/documentation/migration.md
@@ -352,18 +352,18 @@ dataset reads are deterministic as well.
 #### Specify a task (e.g. classification, ranking) instead of a loss (e.g. binary cross-entropy)
 
 ```diff {.bad}
-- model = tf.keras.Sequential()
+- model = tf_keras.Sequential()
 - model.add(Dense(64, activation=relu))
 - model.add(Dense(1)) # One output for binary classification
 
-- model.compile(loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
+- model.compile(loss=tf_keras.losses.BinaryCrossentropy(from_logits=True),
 -               optimizer='adam',
 -               metrics=['accuracy'])
 ```
 
 ```diff {.good}
 # The loss is automatically determined from the task.
-+ model = tfdf.keras.GradientBoostedTreesModel(task=tf.keras.Task.CLASSIFICATION)
++ model = tfdf.keras.GradientBoostedTreesModel(task=tf_keras.Task.CLASSIFICATION)
 
 # Optional if you want to report the accuracy.
 + model.compile(metrics=['accuracy'])

diff --git a/documentation/tf_df_in_tf_js.md b/documentation/tf_df_in_tf_js.md
@@ -45,9 +45,10 @@ import tensorflow as tf
 import tensorflow_decision_forests as tfdf
 import tensorflowjs as tfjs
 from google.colab import files
+import tf_keras
 
 # Load the model with Keras
-model = tf.keras.models.load_model("/tmp/my_saved_model/")
+model = tf_keras.models.load_model("/tmp/my_saved_model/")
 
 # Convert the keras model to TensorFlow.js
 tfjs.converters.tf_saved_model_conversion_v2.convert_keras_model_to_graph_model(model, "./tfjs_model")

diff --git a/documentation/tutorials/advanced_colab.ipynb b/documentation/tutorials/advanced_colab.ipynb
@@ -639,7 +639,7 @@
       },
       "outputs": [],
       "source": [
-        "manual_model = tf.keras.models.load_model(\"/tmp/manual_model\")"
+        "manual_model = tf_keras.models.load_model(\"/tmp/manual_model\")"
       ]
     },
     {

diff --git a/documentation/tutorials/beginner_colab.ipynb b/documentation/tutorials/beginner_colab.ipynb
@@ -1078,16 +1078,16 @@
       "source": [
         "%set_cell_height 300\n",
         "\n",
-        "body_mass_g = tf.keras.layers.Input(shape=(1,), name=\"body_mass_g\")\n",
+        "body_mass_g = tf_keras.layers.Input(shape=(1,), name=\"body_mass_g\")\n",
         "body_mass_kg = body_mass_g / 1000.0\n",
         "\n",
-        "bill_length_mm = tf.keras.layers.Input(shape=(1,), name=\"bill_length_mm\")\n",
+        "bill_length_mm = tf_keras.layers.Input(shape=(1,), name=\"bill_length_mm\")\n",
         "\n",
         "raw_inputs = {\"body_mass_g\": body_mass_g, \"bill_length_mm\": bill_length_mm}\n",
         "processed_inputs = {\"body_mass_kg\": body_mass_kg, \"bill_length_mm\": bill_length_mm}\n",
         "\n",
         "# \"preprocessor\" contains the preprocessing logic.\n",
-        "preprocessor = tf.keras.Model(inputs=raw_inputs, outputs=processed_inputs)\n",
+        "preprocessor = tf_keras.Model(inputs=raw_inputs, outputs=processed_inputs)\n",
         "\n",
         "# \"model_4\" contains both the pre-processing logic and the decision forest.\n",
         "model_4 = tfdf.keras.RandomForestModel(preprocessing=preprocessor)\n",
@@ -1122,7 +1122,7 @@
         "    tf.feature_column.numeric_column(\"bill_length_mm\"),\n",
         "]\n",
         "\n",
-        "preprocessing = tf.keras.layers.DenseFeatures(feature_columns)\n",
+        "preprocessing = tf_keras.layers.DenseFeatures(feature_columns)\n",
         "\n",
         "model_5 = tfdf.keras.RandomForestModel(preprocessing=preprocessing)\n",
         "model_5.fit(train_ds)"

diff --git a/documentation/tutorials/intermediate_colab.ipynb b/documentation/tutorials/intermediate_colab.ipynb
@@ -421,12 +421,12 @@
         "hub_url = \"https://tfhub.dev/google/universal-sentence-encoder/4\"\n",
         "embedding = hub.KerasLayer(hub_url)\n",
         "\n",
-        "sentence = tf.keras.layers.Input(shape=(), name=\"sentence\", dtype=tf.string)\n",
+        "sentence = tf_keras.layers.Input(shape=(), name=\"sentence\", dtype=tf.string)\n",
         "embedded_sentence = embedding(sentence)\n",
         "\n",
         "raw_inputs = {\"sentence\": sentence}\n",
         "processed_inputs = {\"embedded_sentence\": embedded_sentence}\n",
-        "preprocessor = tf.keras.Model(inputs=raw_inputs, outputs=processed_inputs)\n",
+        "preprocessor = tf_keras.Model(inputs=raw_inputs, outputs=processed_inputs)\n",
         "\n",
         "model_2 = tfdf.keras.RandomForestModel(\n",
         "    preprocessing=preprocessor,\n",
@@ -621,8 +621,8 @@
       },
       "outputs": [],
       "source": [
-        "input_1 = tf.keras.Input(shape=(1,), name=\"bill_length_mm\", dtype=\"float\")\n",
+        "input_1 = tf_keras.Input(shape=(1,), name=\"bill_length_mm\", dtype=\"float\")\n",
-        "input_2 = tf.keras.Input(shape=(1,), name=\"island\", dtype=\"string\")\n",
+        "input_2 = tf_keras.Input(shape=(1,), name=\"island\", dtype=\"string\")\n",
         "\n",
         "nn_raw_inputs = [input_1, input_2]"
       ]
@@ -645,9 +645,9 @@
       "outputs": [],
       "source": [
         "# Normalization.\n",
-        "Normalization = tf.keras.layers.Normalization\n",
+        "Normalization = tf_keras.layers.Normalization\n",
-        "CategoryEncoding = tf.keras.layers.CategoryEncoding\n",
+        "CategoryEncoding = tf_keras.layers.CategoryEncoding\n",
-        "StringLookup = tf.keras.layers.StringLookup\n",
+        "StringLookup = tf_keras.layers.StringLookup\n",
         "\n",
         "values = train_ds_pd[\"bill_length_mm\"].values[:, tf.newaxis]\n",
         "input_1_normalizer = Normalization()\n",
@@ -682,15 +682,15 @@
       },
       "outputs": [],
       "source": [
-        "y = tf.keras.layers.Concatenate()(nn_processed_inputs)\n",
+        "y = tf_keras.layers.Concatenate()(nn_processed_inputs)\n",
-        "y = tf.keras.layers.Dense(16, activation=tf.nn.relu6)(y)\n",
+        "y = tf_keras.layers.Dense(16, activation=tf.nn.relu6)(y)\n",
-        "last_layer = tf.keras.layers.Dense(8, activation=tf.nn.relu, name=\"last\")(y)\n",
+        "last_layer = tf_keras.layers.Dense(8, activation=tf.nn.relu, name=\"last\")(y)\n",
         "\n",
         "# \"3\" for the three label classes. If it were a binary classification, the\n",
         "# output dim would be 1.\n",
-        "classification_output = tf.keras.layers.Dense(3)(y)\n",
+        "classification_output = tf_keras.layers.Dense(3)(y)\n",
         "\n",
-        "nn_model = tf.keras.models.Model(nn_raw_inputs, classification_output)"
+        "nn_model = tf_keras.models.Model(nn_raw_inputs, classification_output)"
       ]
     },
     {
@@ -714,7 +714,7 @@
       "source": [
         "# To reduce the risk of mistakes, group both the decision forest and the\n",
         "# neural network in a single keras model.\n",
-        "nn_without_head = tf.keras.models.Model(inputs=nn_model.inputs, outputs=last_layer)\n",
+        "nn_without_head = tf_keras.models.Model(inputs=nn_model.inputs, outputs=last_layer)\n",
         "df_and_nn_model = tfdf.keras.RandomForestModel(preprocessing=nn_without_head)"
       ]
     },
@@ -740,8 +740,8 @@
         "%set_cell_height 300\n",
         "\n",
         "nn_model.compile(\n",
-        "  optimizer=tf.keras.optimizers.Adam(),\n",
+        "  optimizer=tf_keras.optimizers.Adam(),\n",
-        "  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+        "  loss=tf_keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
         "  metrics=[\"accuracy\"])\n",
         "\n",
         "nn_model.fit(x=train_ds, validation_data=test_ds, epochs=10)\n",

diff --git a/documentation/tutorials/model_composition_colab.ipynb b/documentation/tutorials/model_composition_colab.ipynb
@@ -414,25 +414,25 @@
       "outputs": [],
       "source": [
         "# Input features.\n",
-        "raw_features = tf.keras.layers.Input(shape=(num_features,))\n",
+        "raw_features = tf_keras.layers.Input(shape=(num_features,))\n",
         "\n",
         "# Stage 1\n",
         "# =======\n",
         "\n",
         "# Common learnable pre-processing\n",
-        "preprocessor = tf.keras.layers.Dense(10, activation=tf.nn.relu6)\n",
+        "preprocessor = tf_keras.layers.Dense(10, activation=tf.nn.relu6)\n",
         "preprocess_features = preprocessor(raw_features)\n",
         "\n",
         "# Stage 2\n",
         "# =======\n",
         "\n",
         "# Model #1: NN\n",
-        "m1_z1 = tf.keras.layers.Dense(5, activation=tf.nn.relu6)(preprocess_features)\n",
+        "m1_z1 = tf_keras.layers.Dense(5, activation=tf.nn.relu6)(preprocess_features)\n",
-        "m1_pred = tf.keras.layers.Dense(1, activation=tf.nn.sigmoid)(m1_z1)\n",
+        "m1_pred = tf_keras.layers.Dense(1, activation=tf.nn.sigmoid)(m1_z1)\n",
         "\n",
         "# Model #2: NN\n",
-        "m2_z1 = tf.keras.layers.Dense(5, activation=tf.nn.relu6)(preprocess_features)\n",
+        "m2_z1 = tf_keras.layers.Dense(5, activation=tf.nn.relu6)(preprocess_features)\n",
-        "m2_pred = tf.keras.layers.Dense(1, activation=tf.nn.sigmoid)(m2_z1)\n",
+        "m2_pred = tf_keras.layers.Dense(1, activation=tf.nn.sigmoid)(m2_z1)\n",
         "\n",
         "\n",
         "# Model #3: DF\n",
@@ -460,8 +460,8 @@
         "# Keras Models\n",
         "# ============\n",
         "\n",
-        "ensemble_nn_only = tf.keras.models.Model(raw_features, mean_nn_only)\n",
+        "ensemble_nn_only = tf_keras.models.Model(raw_features, mean_nn_only)\n",
-        "ensemble_nn_and_df = tf.keras.models.Model(raw_features, mean_nn_and_df)"
+        "ensemble_nn_and_df = tf_keras.models.Model(raw_features, mean_nn_and_df)"
       ]
     },
     {
@@ -509,8 +509,8 @@
       "source": [
         "%%time\n",
         "ensemble_nn_only.compile(\n",
-        "        optimizer=tf.keras.optimizers.Adam(),\n",
+        "        optimizer=tf_keras.optimizers.Adam(),\n",
-        "        loss=tf.keras.losses.BinaryCrossentropy(),\n",
+        "        loss=tf_keras.losses.BinaryCrossentropy(),\n",
         "        metrics=[\"accuracy\"])\n",
         "\n",
         "ensemble_nn_only.fit(train_dataset, epochs=20, validation_data=test_dataset)"
@@ -610,7 +610,7 @@
       "outputs": [],
       "source": [
         "ensemble_nn_and_df.compile(\n",
-        "    loss=tf.keras.losses.BinaryCrossentropy(), metrics=[\"accuracy\"])\n",
+        "    loss=tf_keras.losses.BinaryCrossentropy(), metrics=[\"accuracy\"])\n",
         "\n",
         "evaluation_nn_and_df = ensemble_nn_and_df.evaluate(\n",
         "    test_dataset, return_dict=True)\n",

diff --git a/documentation/tutorials/predict_colab.ipynb b/documentation/tutorials/predict_colab.ipynb
@@ -93,7 +93,7 @@
         "\n",
         "While abstracted by the Keras API, a model instantiated in Python (e.g., with\n",
         "`tfdf.keras.RandomForestModel()`) and a model loaded from disk (e.g., with\n",
-        "`tf.keras.models.load_model()`) can behave differently. Notably, a Python\n",
+        "`tf_keras.models.load_model()`) can behave differently. Notably, a Python\n",
         "instantiated model automatically applies necessary type conversions. For\n",
         "example, if a `float64` feature is fed to a model expecting a `float32` feature,\n",
         "this conversion is performed implicitly. However, such a conversion is not\n",

diff --git a/documentation/tutorials/ranking_colab.ipynb b/documentation/tutorials/ranking_colab.ipynb
@@ -244,7 +244,7 @@
       },
       "outputs": [],
       "source": [
-        "archive_path = tf.keras.utils.get_file(\"letor.zip\",\n",
+        "archive_path = tf_keras.utils.get_file(\"letor.zip\",\n",
         "  \"https://download.microsoft.com/download/E/7/E/E7EABEF1-4C7B-4E31-ACE5-73927950ED5E/Letor.zip\",\n",
         "  extract=True)\n",
         "\n",

diff --git a/examples/BUILD b/examples/BUILD
@@ -14,6 +14,7 @@ py_binary(
         # pandas dep,
         # TensorFlow Python,
         "//tensorflow_decision_forests",
+        # tf_keras dep,
     ],
 )
 
@@ -28,6 +29,7 @@ py_binary(
         # pandas dep,
         # TensorFlow Python,
         "//tensorflow_decision_forests",
+        # tf_keras dep,
     ],
 )
 
@@ -43,6 +45,7 @@ py_binary(
         # pandas dep,
         # TensorFlow Python,
         "//tensorflow_decision_forests",
+        # tf_keras dep,
     ],
 )
 
@@ -56,5 +59,6 @@ py_binary(
         # absl/logging dep,
         # TensorFlow Python,
         "//tensorflow_decision_forests",
+        # tf_keras dep,
     ],
 )
diff --git a/examples/hyperparameter_optimization.py b/examples/hyperparameter_optimization.py
@@ -36,14 +36,14 @@
 import pandas as pd
 import tensorflow as tf
 import tensorflow_decision_forests as tfdf
-
+import tf_keras
 
 def main(argv):
   if len(argv) > 1:
     raise app.UsageError("Too many command-line arguments.")
 
   # Download the Adult dataset.
-  dataset_path = tf.keras.utils.get_file(
+  dataset_path = tf_keras.utils.get_file(
       "adult.csv",
       "https://raw.githubusercontent.com/google/yggdrasil-decision-forests/"
       "main/yggdrasil_decision_forests/test_data/dataset/adult.csv")

diff --git a/examples/minimal.py b/examples/minimal.py
@@ -30,22 +30,23 @@
 """
 
 from absl import app
-
 import numpy as np
 import pandas as pd
 import tensorflow as tf
 import tensorflow_decision_forests as tfdf
+import tf_keras
 
 
 def main(argv):
   if len(argv) > 1:
     raise app.UsageError("Too many command-line arguments.")
 
   # Download the Adult dataset.
-  dataset_path = tf.keras.utils.get_file(
+  dataset_path = tf_keras.utils.get_file(
       "adult.csv",
       "https://raw.githubusercontent.com/google/yggdrasil-decision-forests/"
-      "main/yggdrasil_decision_forests/test_data/dataset/adult.csv")
+      "main/yggdrasil_decision_forests/test_data/dataset/adult.csv",
+  )
 
   # Load a dataset into a Pandas Dataframe.
   dataset_df = pd.read_csv(dataset_path)  # "df" for Pandas's DataFrame.
@@ -61,8 +62,10 @@ def main(argv):
   test_indices = np.random.rand(len(dataset_df)) < 0.30
   test_ds_pd = dataset_df[test_indices]
   train_ds_pd = dataset_df[~test_indices]
-  print(f"{len(train_ds_pd)} examples in training"
+  print(
-        f", {len(test_ds_pd)} examples for testing.")
+      f"{len(train_ds_pd)} examples in training"
+      f", {len(test_ds_pd)} examples for testing."
+  )
 
   # Converts datasets from Pandas dataframe to TensorFlow dataset format.
   train_ds = tfdf.keras.pd_dataframe_to_tf_dataset(train_ds_pd, label="income")

diff --git a/tensorflow_decision_forests/__init__.py b/tensorflow_decision_forests/__init__.py
@@ -46,7 +46,7 @@
 # ...
 
 # Load a model: it loads as a generic keras model.
-loaded_model = tf.keras.models.load_model("/tmp/my_saved_model")
+loaded_model = tf_keras.models.load_model("/tmp/my_saved_model")
 ```
 
 """

diff --git a/tensorflow_decision_forests/component/builder/builder.py b/tensorflow_decision_forests/component/builder/builder.py
@@ -106,7 +106,7 @@
 builder.close()
 
 # Load and use the model
-model = tf.keras.models.load_model("/path/to/model")
+model = tf_keras.models.load_model("/path/to/model")
 predictions = model.predict(...)
 ```
 """