{ "cells": [ { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "This is a companion notebook for the book [Deep Learning with Python, Third Edition](https://www.manning.com/books/deep-learning-with-python-third-edition). For readability, it only contains runnable code blocks and section titles, and omits everything else in the book: text paragraphs, figures, and pseudocode.\n\n**If you want to be able to follow what's going on, I recommend reading the notebook side by side with your copy of the book.**\n\nThe book's contents are available online at [deeplearningwithpython.io](https://deeplearningwithpython.io)." ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "!pip install keras keras-hub --upgrade -q" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "import os\n", "os.environ[\"KERAS_BACKEND\"] = \"jax\"" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "cellView": "form", "colab_type": "code" }, "outputs": [], "source": [ "# @title\n", "import os\n", "from IPython.core.magic import register_cell_magic\n", "\n", "@register_cell_magic\n", "def backend(line, cell):\n", " current, required = os.environ.get(\"KERAS_BACKEND\", \"\"), line.split()[-1]\n", " if current == required:\n", " get_ipython().run_cell(cell)\n", " else:\n", " print(\n", " f\"This cell requires the {required} backend. To run it, change KERAS_BACKEND to \"\n", " f\"\\\"{required}\\\" at the top of the notebook, restart the runtime, and rerun the notebook.\"\n", " )" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "## A deep dive on Keras" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "### A spectrum of workflows" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "### Different ways to build Keras models" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### The Sequential model" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "import keras\n", "from keras import layers\n", "\n", "model = keras.Sequential(\n", " [\n", " layers.Dense(64, activation=\"relu\"),\n", " layers.Dense(10, activation=\"softmax\"),\n", " ]\n", ")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model = keras.Sequential()\n", "model.add(layers.Dense(64, activation=\"relu\"))\n", "model.add(layers.Dense(10, activation=\"softmax\"))" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model.weights" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model.build(input_shape=(None, 3))\n", "model.weights" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model.summary(line_length=80)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model = keras.Sequential(name=\"my_example_model\")\n", "model.add(layers.Dense(64, activation=\"relu\", name=\"my_first_layer\"))\n", "model.add(layers.Dense(10, activation=\"softmax\", name=\"my_last_layer\"))\n", "model.build((None, 3))\n", "model.summary(line_length=80)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model = keras.Sequential()\n", "model.add(keras.Input(shape=(3,)))\n", "model.add(layers.Dense(64, activation=\"relu\"))" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model.summary(line_length=80)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model.add(layers.Dense(10, activation=\"softmax\"))\n", "model.summary(line_length=80)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### The Functional API" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### A simple example" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "inputs = keras.Input(shape=(3,), name=\"my_input\")\n", "features = layers.Dense(64, activation=\"relu\")(inputs)\n", "outputs = layers.Dense(10, activation=\"softmax\")(features)\n", "model = keras.Model(inputs=inputs, outputs=outputs, name=\"my_functional_model\")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "inputs = keras.Input(shape=(3,), name=\"my_input\")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "inputs.shape" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "inputs.dtype" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "features = layers.Dense(64, activation=\"relu\")(inputs)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "features.shape" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "outputs = layers.Dense(10, activation=\"softmax\")(features)\n", "model = keras.Model(inputs=inputs, outputs=outputs, name=\"my_functional_model\")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model.summary(line_length=80)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### Multi-input, multi-output models" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "vocabulary_size = 10000\n", "num_tags = 100\n", "num_departments = 4\n", "\n", "title = keras.Input(shape=(vocabulary_size,), name=\"title\")\n", "text_body = keras.Input(shape=(vocabulary_size,), name=\"text_body\")\n", "tags = keras.Input(shape=(num_tags,), name=\"tags\")\n", "\n", "features = layers.Concatenate()([title, text_body, tags])\n", "features = layers.Dense(64, activation=\"relu\", name=\"dense_features\")(features)\n", "\n", "priority = layers.Dense(1, activation=\"sigmoid\", name=\"priority\")(features)\n", "department = layers.Dense(\n", " num_departments, activation=\"softmax\", name=\"department\"\n", ")(features)\n", "\n", "model = keras.Model(\n", " inputs=[title, text_body, tags],\n", " outputs=[priority, department],\n", ")" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### Training a multi-input, multi-output model" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "import numpy as np\n", "\n", "num_samples = 1280\n", "\n", "title_data = np.random.randint(0, 2, size=(num_samples, vocabulary_size))\n", "text_body_data = np.random.randint(0, 2, size=(num_samples, vocabulary_size))\n", "tags_data = np.random.randint(0, 2, size=(num_samples, num_tags))\n", "\n", "priority_data = np.random.random(size=(num_samples, 1))\n", "department_data = np.random.randint(0, num_departments, size=(num_samples, 1))\n", "\n", "model.compile(\n", " optimizer=\"adam\",\n", " loss=[\"mean_squared_error\", \"sparse_categorical_crossentropy\"],\n", " metrics=[[\"mean_absolute_error\"], [\"accuracy\"]],\n", ")\n", "model.fit(\n", " [title_data, text_body_data, tags_data],\n", " [priority_data, department_data],\n", " epochs=1,\n", ")\n", "model.evaluate(\n", " [title_data, text_body_data, tags_data], [priority_data, department_data]\n", ")\n", "priority_preds, department_preds = model.predict(\n", " [title_data, text_body_data, tags_data]\n", ")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model.compile(\n", " optimizer=\"adam\",\n", " loss={\n", " \"priority\": \"mean_squared_error\",\n", " \"department\": \"sparse_categorical_crossentropy\",\n", " },\n", " metrics={\n", " \"priority\": [\"mean_absolute_error\"],\n", " \"department\": [\"accuracy\"],\n", " },\n", ")\n", "model.fit(\n", " {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data},\n", " {\"priority\": priority_data, \"department\": department_data},\n", " epochs=1,\n", ")\n", "model.evaluate(\n", " {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data},\n", " {\"priority\": priority_data, \"department\": department_data},\n", ")\n", "priority_preds, department_preds = model.predict(\n", " {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data}\n", ")" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### The power of the Functional API: Access to layer connectivity" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "###### Plotting layer connectivity" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "keras.utils.plot_model(model, \"ticket_classifier.png\")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "keras.utils.plot_model(\n", " model,\n", " \"ticket_classifier_with_shape_info.png\",\n", " show_shapes=True,\n", " show_layer_names=True,\n", ")" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "###### Feature extraction with a Functional model" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model.layers" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model.layers[3].input" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model.layers[3].output" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "features = model.layers[4].output\n", "difficulty = layers.Dense(3, activation=\"softmax\", name=\"difficulty\")(features)\n", "\n", "new_model = keras.Model(\n", " inputs=[title, text_body, tags], outputs=[priority, department, difficulty]\n", ")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "keras.utils.plot_model(\n", " new_model,\n", " \"updated_ticket_classifier.png\",\n", " show_shapes=True,\n", " show_layer_names=True,\n", ")" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Subclassing the Model class" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### Rewriting our previous example as a subclassed model" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "class CustomerTicketModel(keras.Model):\n", " def __init__(self, num_departments):\n", " super().__init__()\n", " self.concat_layer = layers.Concatenate()\n", " self.mixing_layer = layers.Dense(64, activation=\"relu\")\n", " self.priority_scorer = layers.Dense(1, activation=\"sigmoid\")\n", " self.department_classifier = layers.Dense(\n", " num_departments, activation=\"softmax\"\n", " )\n", "\n", " def call(self, inputs):\n", " title = inputs[\"title\"]\n", " text_body = inputs[\"text_body\"]\n", " tags = inputs[\"tags\"]\n", "\n", " features = self.concat_layer([title, text_body, tags])\n", " features = self.mixing_layer(features)\n", " priority = self.priority_scorer(features)\n", " department = self.department_classifier(features)\n", " return priority, department" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model = CustomerTicketModel(num_departments=4)\n", "\n", "priority, department = model(\n", " {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data}\n", ")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model.compile(\n", " optimizer=\"adam\",\n", " loss=[\"mean_squared_error\", \"sparse_categorical_crossentropy\"],\n", " metrics=[[\"mean_absolute_error\"], [\"accuracy\"]],\n", ")\n", "model.fit(\n", " {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data},\n", " [priority_data, department_data],\n", " epochs=1,\n", ")\n", "model.evaluate(\n", " {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data},\n", " [priority_data, department_data],\n", ")\n", "priority_preds, department_preds = model.predict(\n", " {\"title\": title_data, \"text_body\": text_body_data, \"tags\": tags_data}\n", ")" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### Beware: What subclassed models don't support" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Mixing and matching different components" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "class Classifier(keras.Model):\n", " def __init__(self, num_classes=2):\n", " super().__init__()\n", " if num_classes == 2:\n", " num_units = 1\n", " activation = \"sigmoid\"\n", " else:\n", " num_units = num_classes\n", " activation = \"softmax\"\n", " self.dense = layers.Dense(num_units, activation=activation)\n", "\n", " def call(self, inputs):\n", " return self.dense(inputs)\n", "\n", "inputs = keras.Input(shape=(3,))\n", "features = layers.Dense(64, activation=\"relu\")(inputs)\n", "outputs = Classifier(num_classes=10)(features)\n", "model = keras.Model(inputs=inputs, outputs=outputs)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "inputs = keras.Input(shape=(64,))\n", "outputs = layers.Dense(1, activation=\"sigmoid\")(inputs)\n", "binary_classifier = keras.Model(inputs=inputs, outputs=outputs)\n", "\n", "class MyModel(keras.Model):\n", " def __init__(self, num_classes=2):\n", " super().__init__()\n", " self.dense = layers.Dense(64, activation=\"relu\")\n", " self.classifier = binary_classifier\n", "\n", " def call(self, inputs):\n", " features = self.dense(inputs)\n", " return self.classifier(features)\n", "\n", "model = MyModel()" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Remember: Use the right tool for the job" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "### Using built-in training and evaluation loops" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "from keras.datasets import mnist\n", "\n", "def get_mnist_model():\n", " inputs = keras.Input(shape=(28 * 28,))\n", " features = layers.Dense(512, activation=\"relu\")(inputs)\n", " features = layers.Dropout(0.5)(features)\n", " outputs = layers.Dense(10, activation=\"softmax\")(features)\n", " model = keras.Model(inputs, outputs)\n", " return model\n", "\n", "(images, labels), (test_images, test_labels) = mnist.load_data()\n", "images = images.reshape((60000, 28 * 28)).astype(\"float32\") / 255\n", "test_images = test_images.reshape((10000, 28 * 28)).astype(\"float32\") / 255\n", "train_images, val_images = images[10000:], images[:10000]\n", "train_labels, val_labels = labels[10000:], labels[:10000]\n", "\n", "model = get_mnist_model()\n", "model.compile(\n", " optimizer=\"adam\",\n", " loss=\"sparse_categorical_crossentropy\",\n", " metrics=[\"accuracy\"],\n", ")\n", "model.fit(\n", " train_images,\n", " train_labels,\n", " epochs=3,\n", " validation_data=(val_images, val_labels),\n", ")\n", "test_metrics = model.evaluate(test_images, test_labels)\n", "predictions = model.predict(test_images)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Writing your own metrics" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "from keras import ops\n", "\n", "class RootMeanSquaredError(keras.metrics.Metric):\n", " def __init__(self, name=\"rmse\", **kwargs):\n", " super().__init__(name=name, **kwargs)\n", " self.mse_sum = self.add_weight(name=\"mse_sum\", initializer=\"zeros\")\n", " self.total_samples = self.add_weight(\n", " name=\"total_samples\", initializer=\"zeros\"\n", " )\n", "\n", " def update_state(self, y_true, y_pred, sample_weight=None):\n", " y_true = ops.one_hot(y_true, num_classes=ops.shape(y_pred)[1])\n", " mse = ops.sum(ops.square(y_true - y_pred))\n", " self.mse_sum.assign_add(mse)\n", " num_samples = ops.shape(y_pred)[0]\n", " self.total_samples.assign_add(num_samples)\n", "\n", " def result(self):\n", " return ops.sqrt(self.mse_sum / self.total_samples)\n", "\n", " def reset_state(self):\n", " self.mse_sum.assign(0.)\n", " self.total_samples.assign(0.)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model = get_mnist_model()\n", "model.compile(\n", " optimizer=\"adam\",\n", " loss=\"sparse_categorical_crossentropy\",\n", " metrics=[\"accuracy\", RootMeanSquaredError()],\n", ")\n", "model.fit(\n", " train_images,\n", " train_labels,\n", " epochs=3,\n", " validation_data=(val_images, val_labels),\n", ")\n", "test_metrics = model.evaluate(test_images, test_labels)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Using callbacks" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### The EarlyStopping and ModelCheckpoint callbacks" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "callbacks_list = [\n", " keras.callbacks.EarlyStopping(\n", " monitor=\"accuracy\",\n", " patience=1,\n", " ),\n", " keras.callbacks.ModelCheckpoint(\n", " filepath=\"checkpoint_path.keras\",\n", " monitor=\"val_loss\",\n", " save_best_only=True,\n", " ),\n", "]\n", "model = get_mnist_model()\n", "model.compile(\n", " optimizer=\"adam\",\n", " loss=\"sparse_categorical_crossentropy\",\n", " metrics=[\"accuracy\"],\n", ")\n", "model.fit(\n", " train_images,\n", " train_labels,\n", " epochs=10,\n", " callbacks=callbacks_list,\n", " validation_data=(val_images, val_labels),\n", ")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model = keras.models.load_model(\"checkpoint_path.keras\")" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Writing your own callbacks" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "from matplotlib import pyplot as plt\n", "\n", "class LossHistory(keras.callbacks.Callback):\n", " def on_train_begin(self, logs):\n", " self.per_batch_losses = []\n", "\n", " def on_batch_end(self, batch, logs):\n", " self.per_batch_losses.append(logs.get(\"loss\"))\n", "\n", " def on_epoch_end(self, epoch, logs):\n", " plt.clf()\n", " plt.plot(\n", " range(len(self.per_batch_losses)),\n", " self.per_batch_losses,\n", " label=\"Training loss for each batch\",\n", " )\n", " plt.xlabel(f\"Batch (epoch {epoch})\")\n", " plt.ylabel(\"Loss\")\n", " plt.legend()\n", " plt.savefig(f\"plot_at_epoch_{epoch}\", dpi=300)\n", " self.per_batch_losses = []" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model = get_mnist_model()\n", "model.compile(\n", " optimizer=\"adam\",\n", " loss=\"sparse_categorical_crossentropy\",\n", " metrics=[\"accuracy\"],\n", ")\n", "model.fit(\n", " train_images,\n", " train_labels,\n", " epochs=10,\n", " callbacks=[LossHistory()],\n", " validation_data=(val_images, val_labels),\n", ")" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Monitoring and visualization with TensorBoard" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "model = get_mnist_model()\n", "model.compile(\n", " optimizer=\"adam\",\n", " loss=\"sparse_categorical_crossentropy\",\n", " metrics=[\"accuracy\"],\n", ")\n", "\n", "tensorboard = keras.callbacks.TensorBoard(\n", " log_dir=\"/full_path_to_your_log_dir\",\n", ")\n", "model.fit(\n", " train_images,\n", " train_labels,\n", " epochs=10,\n", " validation_data=(val_images, val_labels),\n", " callbacks=[tensorboard],\n", ")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%load_ext tensorboard\n", "%tensorboard --logdir /full_path_to_your_log_dir" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "### Writing your own training and evaluation loops" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Training vs. inference" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Writing custom training step functions" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### A TensorFlow training step function" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend tensorflow\n", "import tensorflow as tf\n", "\n", "model = get_mnist_model()\n", "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n", "optimizer = keras.optimizers.Adam()\n", "\n", "def train_step(inputs, targets):\n", " with tf.GradientTape() as tape:\n", " predictions = model(inputs, training=True)\n", " loss = loss_fn(targets, predictions)\n", " gradients = tape.gradient(loss, model.trainable_weights)\n", " optimizer.apply(gradients, model.trainable_weights)\n", " return loss" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend tensorflow\n", "batch_size = 32\n", "inputs = train_images[:batch_size]\n", "targets = train_labels[:batch_size]\n", "loss = train_step(inputs, targets)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### A PyTorch training step function" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend torch\n", "import torch\n", "\n", "model = get_mnist_model()\n", "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n", "optimizer = keras.optimizers.Adam()\n", "\n", "def train_step(inputs, targets):\n", " predictions = model(inputs, training=True)\n", " loss = loss_fn(targets, predictions)\n", " loss.backward()\n", " gradients = [weight.value.grad for weight in model.trainable_weights]\n", " with torch.no_grad():\n", " optimizer.apply(gradients, model.trainable_weights)\n", " model.zero_grad()\n", " return loss" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend torch\n", "batch_size = 32\n", "inputs = train_images[:batch_size]\n", "targets = train_labels[:batch_size]\n", "loss = train_step(inputs, targets)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### A JAX training step function" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend jax\n", "model = get_mnist_model()\n", "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n", "\n", "def compute_loss_and_updates(\n", " trainable_variables, non_trainable_variables, inputs, targets\n", "):\n", " outputs, non_trainable_variables = model.stateless_call(\n", " trainable_variables, non_trainable_variables, inputs, training=True\n", " )\n", " loss = loss_fn(targets, outputs)\n", " return loss, non_trainable_variables" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend jax\n", "import jax\n", "\n", "grad_fn = jax.value_and_grad(compute_loss_and_updates, has_aux=True)" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend jax\n", "optimizer = keras.optimizers.Adam()\n", "optimizer.build(model.trainable_variables)\n", "\n", "def train_step(state, inputs, targets):\n", " (trainable_variables, non_trainable_variables, optimizer_variables) = state\n", " (loss, non_trainable_variables), grads = grad_fn(\n", " trainable_variables, non_trainable_variables, inputs, targets\n", " )\n", " trainable_variables, optimizer_variables = optimizer.stateless_apply(\n", " optimizer_variables, grads, trainable_variables\n", " )\n", " return loss, (\n", " trainable_variables,\n", " non_trainable_variables,\n", " optimizer_variables,\n", " )" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend jax\n", "batch_size = 32\n", "inputs = train_images[:batch_size]\n", "targets = train_labels[:batch_size]\n", "\n", "trainable_variables = [v.value for v in model.trainable_variables]\n", "non_trainable_variables = [v.value for v in model.non_trainable_variables]\n", "optimizer_variables = [v.value for v in optimizer.variables]\n", "\n", "state = (trainable_variables, non_trainable_variables, optimizer_variables)\n", "loss, state = train_step(state, inputs, targets)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Low-level usage of metrics" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "from keras import ops\n", "\n", "metric = keras.metrics.SparseCategoricalAccuracy()\n", "targets = ops.array([0, 1, 2])\n", "predictions = ops.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])\n", "metric.update_state(targets, predictions)\n", "current_result = metric.result()\n", "print(f\"result: {current_result:.2f}\")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "values = ops.array([0, 1, 2, 3, 4])\n", "mean_tracker = keras.metrics.Mean()\n", "for value in values:\n", " mean_tracker.update_state(value)\n", "print(f\"Mean of values: {mean_tracker.result():.2f}\")" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "metric = keras.metrics.SparseCategoricalAccuracy()\n", "targets = ops.array([0, 1, 2])\n", "predictions = ops.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])\n", "\n", "metric_variables = metric.variables\n", "metric_variables = metric.stateless_update_state(\n", " metric_variables, targets, predictions\n", ")\n", "current_result = metric.stateless_result(metric_variables)\n", "print(f\"result: {current_result:.2f}\")\n", "\n", "metric_variables = metric.stateless_reset_state()" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Using fit() with a custom training loop" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### Customizing fit() with TensorFlow" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend tensorflow\n", "import keras\n", "from keras import layers\n", "\n", "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n", "loss_tracker = keras.metrics.Mean(name=\"loss\")\n", "\n", "class CustomModel(keras.Model):\n", " def train_step(self, data):\n", " inputs, targets = data\n", " with tf.GradientTape() as tape:\n", " predictions = self(inputs, training=True)\n", " loss = loss_fn(targets, predictions)\n", " gradients = tape.gradient(loss, self.trainable_weights)\n", " self.optimizer.apply(gradients, self.trainable_weights)\n", "\n", " loss_tracker.update_state(loss)\n", " return {\"loss\": loss_tracker.result()}\n", "\n", " @property\n", " def metrics(self):\n", " return [loss_tracker]" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend tensorflow\n", "def get_custom_model():\n", " inputs = keras.Input(shape=(28 * 28,))\n", " features = layers.Dense(512, activation=\"relu\")(inputs)\n", " features = layers.Dropout(0.5)(features)\n", " outputs = layers.Dense(10, activation=\"softmax\")(features)\n", " model = CustomModel(inputs, outputs)\n", " model.compile(optimizer=keras.optimizers.Adam())\n", " return model" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend tensorflow\n", "model = get_custom_model()\n", "model.fit(train_images, train_labels, epochs=3)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### Customizing fit() with PyTorch" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend torch\n", "import keras\n", "from keras import layers\n", "\n", "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n", "loss_tracker = keras.metrics.Mean(name=\"loss\")\n", "\n", "class CustomModel(keras.Model):\n", " def train_step(self, data):\n", " inputs, targets = data\n", " predictions = self(inputs, training=True)\n", " loss = loss_fn(targets, predictions)\n", "\n", " loss.backward()\n", " trainable_weights = [v for v in self.trainable_weights]\n", " gradients = [v.value.grad for v in trainable_weights]\n", "\n", " with torch.no_grad():\n", " self.optimizer.apply(gradients, trainable_weights)\n", "\n", " loss_tracker.update_state(loss)\n", " return {\"loss\": loss_tracker.result()}\n", "\n", " @property\n", " def metrics(self):\n", " return [loss_tracker]" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend torch\n", "def get_custom_model():\n", " inputs = keras.Input(shape=(28 * 28,))\n", " features = layers.Dense(512, activation=\"relu\")(inputs)\n", " features = layers.Dropout(0.5)(features)\n", " outputs = layers.Dense(10, activation=\"softmax\")(features)\n", " model = CustomModel(inputs, outputs)\n", " model.compile(optimizer=keras.optimizers.Adam())\n", " return model" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend torch\n", "model = get_custom_model()\n", "model.fit(train_images, train_labels, epochs=3)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### Customizing fit() with JAX" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend jax\n", "import keras\n", "from keras import layers\n", "\n", "loss_fn = keras.losses.SparseCategoricalCrossentropy()\n", "\n", "class CustomModel(keras.Model):\n", " def compute_loss_and_updates(\n", " self,\n", " trainable_variables,\n", " non_trainable_variables,\n", " inputs,\n", " targets,\n", " training=False,\n", " ):\n", " predictions, non_trainable_variables = self.stateless_call(\n", " trainable_variables,\n", " non_trainable_variables,\n", " inputs,\n", " training=training,\n", " )\n", " loss = loss_fn(targets, predictions)\n", " return loss, non_trainable_variables\n", "\n", " def train_step(self, state, data):\n", " (\n", " trainable_variables,\n", " non_trainable_variables,\n", " optimizer_variables,\n", " metrics_variables,\n", " ) = state\n", " inputs, targets = data\n", "\n", " grad_fn = jax.value_and_grad(\n", " self.compute_loss_and_updates, has_aux=True\n", " )\n", "\n", " (loss, non_trainable_variables), grads = grad_fn(\n", " trainable_variables,\n", " non_trainable_variables,\n", " inputs,\n", " targets,\n", " training=True,\n", " )\n", "\n", " (\n", " trainable_variables,\n", " optimizer_variables,\n", " ) = self.optimizer.stateless_apply(\n", " optimizer_variables, grads, trainable_variables\n", " )\n", "\n", " logs = {\"loss\": loss}\n", " state = (\n", " trainable_variables,\n", " non_trainable_variables,\n", " optimizer_variables,\n", " metrics_variables,\n", " )\n", " return logs, state" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend jax\n", "def get_custom_model():\n", " inputs = keras.Input(shape=(28 * 28,))\n", " features = layers.Dense(512, activation=\"relu\")(inputs)\n", " features = layers.Dropout(0.5)(features)\n", " outputs = layers.Dense(10, activation=\"softmax\")(features)\n", " model = CustomModel(inputs, outputs)\n", " model.compile(optimizer=keras.optimizers.Adam())\n", " return model" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend jax\n", "model = get_custom_model()\n", "model.fit(train_images, train_labels, epochs=3)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "#### Handling metrics in a custom train_step()" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### train_step() metrics handling with TensorFlow" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend tensorflow\n", "import keras\n", "from keras import layers\n", "\n", "class CustomModel(keras.Model):\n", " def train_step(self, data):\n", " inputs, targets = data\n", " with tf.GradientTape() as tape:\n", " predictions = self(inputs, training=True)\n", " loss = self.compute_loss(y=targets, y_pred=predictions)\n", "\n", " gradients = tape.gradient(loss, self.trainable_weights)\n", " self.optimizer.apply(gradients, self.trainable_weights)\n", "\n", " for metric in self.metrics:\n", " if metric.name == \"loss\":\n", " metric.update_state(loss)\n", " else:\n", " metric.update_state(targets, predictions)\n", "\n", " return {m.name: m.result() for m in self.metrics}" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend tensorflow\n", "def get_custom_model():\n", " inputs = keras.Input(shape=(28 * 28,))\n", " features = layers.Dense(512, activation=\"relu\")(inputs)\n", " features = layers.Dropout(0.5)(features)\n", " outputs = layers.Dense(10, activation=\"softmax\")(features)\n", " model = CustomModel(inputs, outputs)\n", " model.compile(\n", " optimizer=keras.optimizers.Adam(),\n", " loss=keras.losses.SparseCategoricalCrossentropy(),\n", " metrics=[keras.metrics.SparseCategoricalAccuracy()],\n", " )\n", " return model\n", "\n", "model = get_custom_model()\n", "model.fit(train_images, train_labels, epochs=3)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### train_step() metrics handling with PyTorch" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend torch\n", "import keras\n", "from keras import layers\n", "\n", "class CustomModel(keras.Model):\n", " def train_step(self, data):\n", " inputs, targets = data\n", " predictions = self(inputs, training=True)\n", " loss = self.compute_loss(y=targets, y_pred=predictions)\n", "\n", " loss.backward()\n", " trainable_weights = [v for v in self.trainable_weights]\n", " gradients = [v.value.grad for v in trainable_weights]\n", "\n", " with torch.no_grad():\n", " self.optimizer.apply(gradients, trainable_weights)\n", "\n", " for metric in self.metrics:\n", " if metric.name == \"loss\":\n", " metric.update_state(loss)\n", " else:\n", " metric.update_state(targets, predictions)\n", "\n", " return {m.name: m.result() for m in self.metrics}" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend torch\n", "def get_custom_model():\n", " inputs = keras.Input(shape=(28 * 28,))\n", " features = layers.Dense(512, activation=\"relu\")(inputs)\n", " features = layers.Dropout(0.5)(features)\n", " outputs = layers.Dense(10, activation=\"softmax\")(features)\n", " model = CustomModel(inputs, outputs)\n", " model.compile(\n", " optimizer=keras.optimizers.Adam(),\n", " loss=keras.losses.SparseCategoricalCrossentropy(),\n", " metrics=[keras.metrics.SparseCategoricalAccuracy()],\n", " )\n", " return model\n", "\n", "model = get_custom_model()\n", "model.fit(train_images, train_labels, epochs=3)" ] }, { "cell_type": "markdown", "metadata": { "colab_type": "text" }, "source": [ "##### train_step() metrics handling with JAX" ] }, { "cell_type": "code", "execution_count": 0, "metadata": { "colab_type": "code" }, "outputs": [], "source": [ "%%backend jax\n", "import keras\n", "from keras import layers\n", "\n", "class CustomModel(keras.Model):\n", " def compute_loss_and_updates(\n", " self,\n", " trainable_variables,\n", " non_trainable_variables,\n", " inputs,\n", " targets,\n", " training=False,\n", " ):\n", " predictions, non_trainable_variables = self.stateless_call(\n", " trainable_variables,\n", " non_trainable_variables,\n", " inputs,\n", " training=training,\n", " )\n", " loss = self.compute_loss(y=targets, y_pred=predictions)\n", " return loss, (predictions, non_trainable_variables)\n", "\n", " def train_step(self, state, data):\n", " (\n", " trainable_variables,\n", " non_trainable_variables,\n", " optimizer_variables,\n", " metrics_variables,\n", " ) = state\n", " inputs, targets = data\n", "\n", " grad_fn = jax.value_and_grad(\n", " self.compute_loss_and_updates, has_aux=True\n", " )\n", "\n", " (loss, (predictions, non_trainable_variables)), grads = grad_fn(\n", " trainable_variables,\n", " non_trainable_variables,\n", " inputs,\n", " targets,\n", " training=True,\n", " )\n", " (\n", " trainable_variables,\n", " optimizer_variables,\n", " ) = self.optimizer.stateless_apply(\n", " optimizer_variables, grads, trainable_variables\n", " )\n", "\n", " new_metrics_vars = []\n", " logs = {}\n", " for metric in self.metrics:\n", " num_prev = len(new_metrics_vars)\n", " num_current = len(metric.variables)\n", " current_vars = metrics_variables[num_prev : num_prev + num_current]\n", " if metric.name == \"loss\":\n", " current_vars = metric.stateless_update_state(current_vars, loss)\n", " else:\n", " current_vars = metric.stateless_update_state(\n", " current_vars, targets, predictions\n", " )\n", " logs[metric.name] = metric.stateless_result(current_vars)\n", " new_metrics_vars += current_vars\n", "\n", " state = (\n", " trainable_variables,\n", " non_trainable_variables,\n", " optimizer_variables,\n", " new_metrics_vars,\n", " )\n", " return logs, state" ] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [], "name": "chapter07_deep-dive-keras", "private_outputs": false, "provenance": [], "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.0" } }, "nbformat": 4, "nbformat_minor": 0 }