Updated masked image modeling example for Keras 3.

hertschuh · Nov 18, 2023 · d78359f · d78359f
1 parent 42d2b1a
commit d78359f
Show file tree

Hide file tree

Showing 29 changed files with 303 additions and 267 deletions.
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_13_1.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_13_1.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_13_2.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_13_2.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_13_3.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_13_3.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_17_0.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_17_0.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_17_1.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_17_1.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_32_0.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_32_0.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_1.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_1.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_11.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_11.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_13.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_13.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_15.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_15.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_17.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_17.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_19.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_19.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_21.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_21.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_23.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_23.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_25.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_25.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_27.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_27.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_29.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_29.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_3.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_3.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_31.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_31.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_33.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_33.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_35.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_35.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_37.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_37.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_39.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_39.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_5.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_5.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_7.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_7.png
diff --git a/examples/vision/img/masked_image_modeling/masked_image_modeling_34_9.png b/examples/vision/img/masked_image_modeling/masked_image_modeling_34_9.png
diff --git a/examples/vision/ipynb/masked_image_modeling.ipynb b/examples/vision/ipynb/masked_image_modeling.ipynb
@@ -59,14 +59,7 @@
     "colab_type": "text"
    },
    "source": [
-    "## Imports\n",
-    "\n",
-    "This example requires TensorFlow Addons, which can be installed using the following\n",
-    "command:\n",
-    "\n",
-    "```shell\n",
-    "pip install -U tensorflow-addons\n",
-    "```"
+    "## Imports"
    ]
   },
   {
@@ -77,10 +70,13 @@
    },
    "outputs": [],
    "source": [
-    "from tensorflow.keras import layers\n",
-    "import tensorflow_addons as tfa\n",
-    "from tensorflow import keras\n",
+    "import os\n",
+    "\n",
+    "os.environ[\"KERAS_BACKEND\"] = \"tensorflow\"\n",
+    "\n",
     "import tensorflow as tf\n",
+    "import keras\n",
+    "from keras import layers\n",
     "\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
@@ -233,7 +229,10 @@
     "\n",
     "def get_test_augmentation_model():\n",
     "    model = keras.Sequential(\n",
-    "        [layers.Rescaling(1 / 255.0), layers.Resizing(IMAGE_SIZE, IMAGE_SIZE),],\n",
+    "        [\n",
+    "            layers.Rescaling(1 / 255.0),\n",
+    "            layers.Resizing(IMAGE_SIZE, IMAGE_SIZE),\n",
+    "        ],\n",
     "        name=\"test_data_augmentation\",\n",
     "    )\n",
     "    return model\n",
@@ -775,7 +774,7 @@
     "        loss_output = tf.gather(decoder_patches, mask_indices, axis=1, batch_dims=1)\n",
     "\n",
     "        # Compute the total loss.\n",
-    "        total_loss = self.compiled_loss(loss_patch, loss_output)\n",
+    "        total_loss = self.compute_loss(y=loss_patch, y_pred=loss_output)\n",
     "\n",
     "        return total_loss, loss_patch, loss_output\n",
     "\n",
@@ -793,21 +792,27 @@
     "        ]\n",
     "        grads = tape.gradient(total_loss, train_vars)\n",
     "        tv_list = []\n",
-    "        for (grad, var) in zip(grads, train_vars):\n",
+    "        for grad, var in zip(grads, train_vars):\n",
     "            for g, v in zip(grad, var):\n",
     "                tv_list.append((g, v))\n",
     "        self.optimizer.apply_gradients(tv_list)\n",
     "\n",
     "        # Report progress.\n",
-    "        self.compiled_metrics.update_state(loss_patch, loss_output)\n",
-    "        return {m.name: m.result() for m in self.metrics}\n",
+    "        results = {}\n",
+    "        for metric in self.metrics:\n",
+    "            metric.update_state(loss_patch, loss_output)\n",
+    "            results[metric.name] = metric.result()\n",
+    "        return results\n",
     "\n",
     "    def test_step(self, images):\n",
     "        total_loss, loss_patch, loss_output = self.calculate_loss(images, test=True)\n",
     "\n",
     "        # Update the trackers.\n",
-    "        self.compiled_metrics.update_state(loss_patch, loss_output)\n",
-    "        return {m.name: m.result() for m in self.metrics}\n",
+    "        results = {}\n",
+    "        for metric in self.metrics:\n",
+    "            metric.update_state(loss_patch, loss_output)\n",
+    "            results[metric.name] = metric.result()\n",
+    "        return results\n",
     ""
    ]
   },
@@ -1022,14 +1027,19 @@
    },
    "outputs": [],
    "source": [
-    "optimizer = tfa.optimizers.AdamW(learning_rate=scheduled_lrs, weight_decay=WEIGHT_DECAY)\n",
+    "optimizer = keras.optimizers.AdamW(\n",
+    "    learning_rate=scheduled_lrs, weight_decay=WEIGHT_DECAY\n",
+    ")\n",
     "\n",
     "# Compile and pretrain the model.\n",
     "mae_model.compile(\n",
     "    optimizer=optimizer, loss=keras.losses.MeanSquaredError(), metrics=[\"mae\"]\n",
     ")\n",
     "history = mae_model.fit(\n",
-    "    train_ds, epochs=EPOCHS, validation_data=val_ds, callbacks=train_callbacks,\n",
+    "    train_ds,\n",
+    "    epochs=EPOCHS,\n",
+    "    validation_data=val_ds,\n",
+    "    callbacks=train_callbacks,\n",
     ")\n",
     "\n",
     "# Measure its performance.\n",

diff --git a/examples/vision/masked_image_modeling.py b/examples/vision/masked_image_modeling.py
@@ -43,19 +43,14 @@
 
 """
 ## Imports
-
-This example requires TensorFlow Addons, which can be installed using the following
-command:
-
-```shell
-pip install -U tensorflow-addons
-```
 """
+import os
+
+os.environ["KERAS_BACKEND"] = "tensorflow"
 
-from tensorflow.keras import layers
-import tensorflow_addons as tfa
-from tensorflow import keras
 import tensorflow as tf
+import keras
+from keras import layers
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -599,7 +594,7 @@ def calculate_loss(self, images, test=False):
         loss_output = tf.gather(decoder_patches, mask_indices, axis=1, batch_dims=1)
 
         # Compute the total loss.
-        total_loss = self.compiled_loss(loss_patch, loss_output)
+        total_loss = self.compute_loss(y=loss_patch, y_pred=loss_output)
 
         return total_loss, loss_patch, loss_output
 
@@ -623,15 +618,21 @@ def train_step(self, images):
         self.optimizer.apply_gradients(tv_list)
 
         # Report progress.
-        self.compiled_metrics.update_state(loss_patch, loss_output)
-        return {m.name: m.result() for m in self.metrics}
+        results = {}
+        for metric in self.metrics:
+            metric.update_state(loss_patch, loss_output)
+            results[metric.name] = metric.result()
+        return results
 
     def test_step(self, images):
         total_loss, loss_patch, loss_output = self.calculate_loss(images, test=True)
 
         # Update the trackers.
-        self.compiled_metrics.update_state(loss_patch, loss_output)
-        return {m.name: m.result() for m in self.metrics}
+        results = {}
+        for metric in self.metrics:
+            metric.update_state(loss_patch, loss_output)
+            results[metric.name] = metric.result()
+        return results
 
 
 """
@@ -785,7 +786,9 @@ def __call__(self, step):
 ## Model compilation and training
 """
 
-optimizer = tfa.optimizers.AdamW(learning_rate=scheduled_lrs, weight_decay=WEIGHT_DECAY)
+optimizer = keras.optimizers.AdamW(
+    learning_rate=scheduled_lrs, weight_decay=WEIGHT_DECAY
+)
 
 # Compile and pretrain the model.
 mae_model.compile(