update tfnp example

examples/keras_recipes/ipynb/tensorflow_numpy_models.ipynb

@@ -41,8 +41,7 @@
  "colab_type": "text"
  },
  "source": [
- "## Setup\n",
- "TensorFlow NumPy requires TensorFlow 2.5 or later."
+ "## Setup"
  ]
  },
  {
@@ -53,32 +52,14 @@
  },
  "outputs": [],
  "source": [
+ "import os\n",
+ "\n",
+ "os.environ[\"KERAS_BACKEND\"] = \"tensorflow\"\n",
+ "\n",
  "import tensorflow as tf\n",
  "import tensorflow.experimental.numpy as tnp\n",
  "import keras\n",
- "import keras.layers as layers\n",
- "import numpy as np"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "colab_type": "text"
- },
- "source": [
- "Optionally, you can call `tnp.experimental_enable_numpy_behavior()` to enable type promotion in TensorFlow.\n",
- "This allows TNP to more closely follow the NumPy standard."
- ]
- },
- {
- "cell_type": "code",
- "execution_count": 0,
- "metadata": {
- "colab_type": "code"
- },
- "outputs": [],
- "source": [
- "tnp.experimental_enable_numpy_behavior()"
+ "from keras import layers"
  ]
  },
  {
@@ -98,16 +79,17 @@
  },
  "outputs": [],
  "source": [
- "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.boston_housing.load_data(\n",
+ "(x_train, y_train), (x_test, y_test) = keras.datasets.boston_housing.load_data(\n",
  " path=\"boston_housing.npz\", test_split=0.2, seed=113\n",
  ")\n",
+ "input_dim = x_train.shape[1]\n",
  "\n",
  "\n",
  "def evaluate_model(model: keras.Model):\n",
- " [loss, percent_error] = model.evaluate(x_test, y_test, verbose=0)\n",
+ " loss, percent_error = model.evaluate(x_test, y_test, verbose=0)\n",
  " print(\"Mean absolute percent error before training: \", percent_error)\n",
  " model.fit(x_train, y_train, epochs=200, verbose=0)\n",
- " [loss, percent_error] = model.evaluate(x_test, y_test, verbose=0)\n",
+ " loss, percent_error = model.evaluate(x_test, y_test, verbose=0)\n",
  " print(\"Mean absolute percent error after training:\", percent_error)\n",
  ""
  ]
@@ -157,16 +139,27 @@
  " for i, block in enumerate(self.blocks):\n",
  " self.block_weights.append(\n",
  " self.add_weight(\n",
- " shape=(current_shape, block), trainable=True, name=f\"block-{i}\"\n",
+ " shape=(current_shape, block),\n",
+ " trainable=True,\n",
+ " name=f\"block-{i}\",\n",
+ " initializer=\"glorot_normal\",\n",
  " )\n",
  " )\n",
  " self.biases.append(\n",
- " self.add_weight(shape=(block,), trainable=True, name=f\"bias-{i}\")\n",
+ " self.add_weight(\n",
+ " shape=(block,),\n",
+ " trainable=True,\n",
+ " name=f\"bias-{i}\",\n",
+ " initializer=\"zeros\",\n",
+ " )\n",
  " )\n",
  " current_shape = block\n",
  "\n",
  " self.linear_layer = self.add_weight(\n",
- " shape=(current_shape, 1), name=\"linear_projector\", trainable=True\n",
+ " shape=(current_shape, 1),\n",
+ " name=\"linear_projector\",\n",
+ " trainable=True,\n",
+ " initializer=\"glorot_normal\",\n",
  " )\n",
  "\n",
  " def call(self, inputs):\n",
@@ -215,7 +208,7 @@
  "colab_type": "text"
  },
  "source": [
- "Great!  Our model seems to be effectively learning to solve the problem at hand.\n",
+ "Great! Our model seems to be effectively learning to solve the problem at hand.\n",
  "\n",
  "We can also write our own custom loss function using TNP."
  ]
@@ -284,7 +277,7 @@
  " self.bias = self.add_weight(\n",
  " name=\"bias\",\n",
  " shape=(self.units,),\n",
- " initializer=\"random_normal\",\n",
+ " initializer=\"zeros\",\n",
  " trainable=True,\n",
  " )\n",
  "\n",
@@ -311,7 +304,7 @@
  " loss=\"mean_squared_error\",\n",
  " metrics=[keras.metrics.MeanAbsolutePercentageError()],\n",
  ")\n",
- "model.build((None, 13,))\n",
+ "model.build((None, input_dim))\n",
  "model.summary()\n",
  "\n",
  "evaluate_model(model)"
@@ -353,7 +346,7 @@
  " loss=\"mean_squared_error\",\n",
  " metrics=[keras.metrics.MeanAbsolutePercentageError()],\n",
  ")\n",
- "model.build((None, 13,))\n",
+ "model.build((None, input_dim))\n",
  "model.summary()\n",
  "\n",
  "evaluate_model(model)"
@@ -406,7 +399,7 @@
  " loss=\"mean_squared_error\",\n",
  " metrics=[keras.metrics.MeanAbsolutePercentageError()],\n",
  " )\n",
- " model.build((None, 13,))\n",
+ " model.build((None, input_dim))\n",
  " model.summary()\n",
  " evaluate_model(model)"
  ]
@@ -456,7 +449,10 @@
  "outputs": [],
  "source": [
  "models = [\n",
- " (TNPForwardFeedRegressionNetwork(blocks=[3, 3]), \"TNPForwardFeedRegressionNetwork\"),\n",
+ " (\n",
+ " TNPForwardFeedRegressionNetwork(blocks=[3, 3]),\n",
+ " \"TNPForwardFeedRegressionNetwork\",\n",
+ " ),\n",
  " (create_layered_tnp_model(), \"layered_tnp_model\"),\n",
  " (create_mixed_model(), \"mixed_model\"),\n",
  "]\n",
@@ -546,4 +542,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
+}