Update knowledge distillation for Keras3

examples/generative/ipynb/pixelcnn.ipynb

@@ -29,9 +29,9 @@
     "probability distribution of later elements. In the following example, images are generated\n",
     "in this fashion, pixel-by-pixel, via a masked convolution kernel that only looks at data\n",
     "from previously generated pixels (origin at the top left) to generate later pixels.\n",
-    "During inference, the output of the network is used as a probability distribution\n",
+    "During inference, the output of the network is used as a probability ditribution\n",
     "from which new pixel values are sampled to generate a new image\n",
-    "(here, with MNIST, the pixel values range from white (0) to black (255).\n"
+    "(here, with MNIST, the pixels values are either black or white)."
    ]
   },
   {
@@ -43,10 +43,10 @@
    "outputs": [],
    "source": [
     "import numpy as np\n",
-    "import tensorflow as tf\n",
-    "from tensorflow import keras\n",
-    "from tensorflow.keras import layers\n",
-    "from tqdm import tqdm\n"
+    "import keras\n",
+    "from keras import layers\n",
+    "from keras import ops\n",
+    "from tqdm import tqdm"
    ]
   },
   {
@@ -55,7 +55,7 @@
     "colab_type": "text"
    },
    "source": [
-    "## Getting the data\n"
+    "## Getting the Data"
    ]
   },
   {
@@ -78,7 +78,7 @@
     "# anything above this value gets rounded up to 1 so that all values are either\n",
     "# 0 or 1\n",
     "data = np.where(data < (0.33 * 256), 0, 1)\n",
-    "data = data.astype(np.float32)\n"
+    "data = data.astype(np.float32)"
    ]
   },
   {
@@ -87,7 +87,7 @@
     "colab_type": "text"
    },
    "source": [
-    "## Create two classes for the requisite Layers for the model\n"
+    "## Create two classes for the requisite Layers for the model"
    ]
   },
   {
@@ -98,6 +98,7 @@
    },
    "outputs": [],
    "source": [
+    "\n",
     "# The first layer is the PixelCNN layer. This layer simply\n",
     "# builds on the 2D convolutional layer, but includes masking.\n",
     "class PixelConvLayer(layers.Layer):\n",
@@ -110,7 +111,7 @@
     "        # Build the conv2d layer to initialize kernel variables\n",
     "        self.conv.build(input_shape)\n",
     "        # Use the initialized kernel to create the mask\n",
-    "        kernel_shape = self.conv.kernel.get_shape()\n",
+    "        kernel_shape = ops.shape(self.conv.kernel)\n",
     "        self.mask = np.zeros(shape=kernel_shape)\n",
     "        self.mask[: kernel_shape[0] // 2, ...] = 1.0\n",
     "        self.mask[kernel_shape[0] // 2, : kernel_shape[1] // 2, ...] = 1.0\n",
@@ -146,7 +147,7 @@
     "        x = self.pixel_conv(x)\n",
     "        x = self.conv2(x)\n",
     "        return keras.layers.add([inputs, x])\n",
-    "\n"
+    ""
    ]
   },
   {
@@ -155,7 +156,7 @@
     "colab_type": "text"
    },
    "source": [
-    "## Build the model based on the original paper\n"
+    "## Build the model based on the original paper"
    ]
   },
   {
@@ -166,7 +167,7 @@
    },
    "outputs": [],
    "source": [
-    "inputs = keras.Input(shape=input_shape)\n",
+    "inputs = keras.Input(shape=input_shape, batch_size=128)\n",
     "x = PixelConvLayer(\n",
     "    mask_type=\"A\", filters=128, kernel_size=7, activation=\"relu\", padding=\"same\"\n",
     ")(inputs)\n",
@@ -195,7 +196,7 @@
     "pixel_cnn.summary()\n",
     "pixel_cnn.fit(\n",
     "    x=data, y=data, batch_size=128, epochs=50, validation_split=0.1, verbose=2\n",
-    ")\n"
+    ")"
    ]
   },
   {
@@ -208,10 +209,7 @@
     "\n",
     "The PixelCNN cannot generate the full image at once. Instead, it must generate each pixel in\n",
     "order, append the last generated pixel to the current image, and feed the image back into the\n",
-    "model to repeat the process.\n",
-    "\n",
-    "You can use the trained model hosted on [Hugging Face Hub](https://huggingface.co/keras-io/pixel-cnn-mnist) ",
-    "and try the demo on [Hugging Face Spaces](https://huggingface.co/spaces/keras-io/pixelcnn-mnist-image-generation)."
+    "model to repeat the process."
    ]
   },
   {
@@ -238,13 +236,13 @@
     "            probs = pixel_cnn.predict(pixels)[:, row, col, channel]\n",
     "            # Use the probabilities to pick pixel values and append the values to the image\n",
     "            # frame.\n",
-    "            pixels[:, row, col, channel] = tf.math.ceil(\n",
-    "                probs - tf.random.uniform(probs.shape)\n",
+    "            pixels[:, row, col, channel] = ops.ceil(\n",
+    "                probs - keras.random.uniform(probs.shape)\n",
     "            )\n",
     "\n",
     "\n",
     "def deprocess_image(x):\n",
-    "    # Stack the single channeled black and white image to RGB values.\n",
+    "    # Stack the single channeled black and white image to rgb values.\n",
     "    x = np.stack((x, x, x), 2)\n",
     "    # Undo preprocessing\n",
     "    x *= 255.0\n",
@@ -262,7 +260,7 @@
     "display(Image(\"generated_image_0.png\"))\n",
     "display(Image(\"generated_image_1.png\"))\n",
     "display(Image(\"generated_image_2.png\"))\n",
-    "display(Image(\"generated_image_3.png\"))\n"
+    "display(Image(\"generated_image_3.png\"))"
    ]
   }
  ],
@@ -295,4 +293,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
+}