Subclass_Model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3efaef4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 11/12/2021-1\n",
    "#https://www.tensorflow.org/guide/keras/customizing_what_happens_in_fit\n",
    "import tensorflow\n",
    "from tensorflow import keras"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "191300f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "class CustomModel(keras.Model):\n",
    "    '''A first simple example '''\n",
    "    def train_step(self, data):\n",
    "        ''' '''\n",
    "        # Unpack the data. Its structure depends on your model and\n",
    "        # on what you pass to `fit()`.\n",
    "        x, y = data\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            y_pred = self(x, training=True)  # Forward pass\n",
    "            # Compute the loss value\n",
    "            # (the loss function is configured in `compile()`)\n",
    "            loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)\n",
    "\n",
    "        # Compute gradients\n",
    "        trainable_vars = self.trainable_variables\n",
    "        gradients = tape.gradient(loss, trainable_vars)\n",
    "        # Update weights\n",
    "        self.optimizer.apply_gradients(zip(gradients, trainable_vars))\n",
    "        # Update metrics (includes the metric that tracks the loss)\n",
    "        self.compiled_metrics.update_state(y, y_pred)\n",
    "        # Return a dict mapping metric names to current value\n",
    "        return {m.name: m.result() for m in self.metrics}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f1c998",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Construct and compile an instance of CustomModel\n",
    "inputs = keras.Input(shape=(32,))\n",
    "outputs = keras.layers.Dense(1)(inputs)\n",
    "model = CustomModel(inputs, outputs)\n",
    "model.compile(optimizer=\"adam\", loss=\"mse\", metrics=[\"mae\"])\n",
    "\n",
    "# Just use `fit` as usual\n",
    "x = np.random.random((1000, 32))\n",
    "y = np.random.random((1000, 1))\n",
    "try: model.fit(x, y, epochs=3)\n",
    "except Exception as err: print(err)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "033c35bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_tracker = keras.metrics.Mean(name=\"loss\")\n",
    "mae_metric = keras.metrics.MeanAbsoluteError(name=\"mae\")\n",
    "\n",
    "\n",
    "class CustomModel(keras.Model):\n",
    "    '''Going lower-level '''\n",
    "    def train_step(self, data):\n",
    "        ''' '''\n",
    "        x, y = data\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            y_pred = self(x, training=True)  # Forward pass\n",
    "            # Compute our own loss\n",
    "            loss = keras.losses.mean_squared_error(y, y_pred)\n",
    "\n",
    "        # Compute gradients\n",
    "        trainable_vars = self.trainable_variables\n",
    "        gradients = tape.gradient(loss, trainable_vars)\n",
    "\n",
    "        # Update weights\n",
    "        self.optimizer.apply_gradients(zip(gradients, trainable_vars))\n",
    "\n",
    "        # Compute our own metrics\n",
    "        loss_tracker.update_state(loss)\n",
    "        mae_metric.update_state(y, y_pred)\n",
    "        return {\"loss\": loss_tracker.result(), \"mae\": mae_metric.result()}\n",
    "\n",
    "    @property\n",
    "    def metrics(self):\n",
    "        # We list our `Metric` objects here so that `reset_states()` can be\n",
    "        # called automatically at the start of each epoch\n",
    "        # or at the start of `evaluate()`.\n",
    "        # If you don't implement this property, you have to call\n",
    "        # `reset_states()` yourself at the time of your choosing.\n",
    "        return [loss_tracker, mae_metric]\n",
    "\n",
    "\n",
    "# Construct an instance of CustomModel\n",
    "inputs = keras.Input(shape=(32,))\n",
    "outputs = keras.layers.Dense(1)(inputs)\n",
    "model = CustomModel(inputs, outputs)\n",
    "\n",
    "# We don't passs a loss or metrics here.\n",
    "model.compile(optimizer=\"adam\")\n",
    "\n",
    "# Just use `fit` as usual -- you can use callbacks, etc.\n",
    "x = np.random.random((1000, 32))\n",
    "y = np.random.random((1000, 1))\n",
    "model.fit(x, y, epochs=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3578006",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel(keras.Model):\n",
    "    '''Supporting sample_weight & class_weight '''\n",
    "    def train_step(self, data):\n",
    "        ''' '''\n",
    "        # Unpack the data. Its structure depends on your model and\n",
    "        # on what you pass to `fit()`.\n",
    "        if len(data) == 3:\n",
    "            x, y, sample_weight = data\n",
    "        else:\n",
    "            sample_weight = None\n",
    "            x, y = data\n",
    "\n",
    "        with tf.GradientTape() as tape:\n",
    "            y_pred = self(x, training=True)  # Forward pass\n",
    "            # Compute the loss value.\n",
    "            # The loss function is configured in `compile()`.\n",
    "            loss = self.compiled_loss(\n",
    "                y,\n",
    "                y_pred,\n",
    "                sample_weight=sample_weight,\n",
    "                regularization_losses=self.losses,\n",
    "            )\n",
    "\n",
    "        # Compute gradients\n",
    "        trainable_vars = self.trainable_variables\n",
    "        gradients = tape.gradient(loss, trainable_vars)\n",
    "\n",
    "        # Update weights\n",
    "        self.optimizer.apply_gradients(zip(gradients, trainable_vars))\n",
    "\n",
    "        # Update the metrics.\n",
    "        # Metrics are configured in `compile()`.\n",
    "        self.compiled_metrics.update_state(y, y_pred, sample_weight=sample_weight)\n",
    "\n",
    "        # Return a dict mapping metric names to current value.\n",
    "        # Note that it will include the loss (tracked in self.metrics).\n",
    "        return {m.name: m.result() for m in self.metrics}\n",
    "\n",
    "\n",
    "# Construct and compile an instance of CustomModel\n",
    "inputs = keras.Input(shape=(32,))\n",
    "outputs = keras.layers.Dense(1)(inputs)\n",
    "model = CustomModel(inputs, outputs)\n",
    "model.compile(optimizer=\"adam\", loss=\"mse\", metrics=[\"mae\"])\n",
    "\n",
    "# You can now use sample_weight argument\n",
    "x = np.random.random((1000, 32))\n",
    "y = np.random.random((1000, 1))\n",
    "sw = np.random.random((1000, 1))\n",
    "model.fit(x, y, sample_weight=sw, epochs=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe06afaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel(keras.Model):\n",
    "    '''Providing your own evaluation step '''\n",
    "    def test_step(self, data):\n",
    "        ''' '''\n",
    "        # Unpack the data\n",
    "        x, y = data\n",
    "        # Compute predictions\n",
    "        y_pred = self(x, training=False)\n",
    "        # Updates the metrics tracking the loss\n",
    "        self.compiled_loss(y, y_pred, regularization_losses=self.losses)\n",
    "        # Update the metrics.\n",
    "        self.compiled_metrics.update_state(y, y_pred)\n",
    "        # Return a dict mapping metric names to current value.\n",
    "        # Note that it will include the loss (tracked in self.metrics).\n",
    "        return {m.name: m.result() for m in self.metrics}\n",
    "\n",
    "\n",
    "# Construct an instance of CustomModel\n",
    "inputs = keras.Input(shape=(32,))\n",
    "outputs = keras.layers.Dense(1)(inputs)\n",
    "model = CustomModel(inputs, outputs)\n",
    "model.compile(loss=\"mse\", metrics=[\"mae\"])\n",
    "\n",
    "# Evaluate with our custom test_step\n",
    "x = np.random.random((1000, 32))\n",
    "y = np.random.random((1000, 1))\n",
    "model.evaluate(x, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60d6f1ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "\n",
    "# Create the discriminator\n",
    "discriminator = keras.Sequential(\n",
    "    [\n",
    "        keras.Input(shape=(28, 28, 1)),\n",
    "        layers.Conv2D(64, (3, 3), strides=(2, 2), padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Conv2D(128, (3, 3), strides=(2, 2), padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.GlobalMaxPooling2D(),\n",
    "        layers.Dense(1),\n",
    "    ],\n",
    "    name=\"discriminator\",\n",
    ")\n",
    "\n",
    "# Create the generator\n",
    "latent_dim = 128\n",
    "generator = keras.Sequential(\n",
    "    [\n",
    "        keras.Input(shape=(latent_dim,)),\n",
    "        # We want to generate 128 coefficients to reshape into a 7x7x128 map\n",
    "        layers.Dense(7 * 7 * 128),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Reshape((7, 7, 128)),\n",
    "        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Conv2DTranspose(128, (4, 4), strides=(2, 2), padding=\"same\"),\n",
    "        layers.LeakyReLU(alpha=0.2),\n",
    "        layers.Conv2D(1, (7, 7), padding=\"same\", activation=\"sigmoid\"),\n",
    "    ],\n",
    "    name=\"generator\",\n",
    ")\n",
    "\n",
    "class GAN(keras.Model):\n",
    "    ''' '''\n",
    "    def __init__(self, discriminator, generator, latent_dim):\n",
    "        super(GAN, self).__init__()\n",
    "        self.discriminator = discriminator\n",
    "        self.generator = generator\n",
    "        self.latent_dim = latent_dim\n",
    "\n",
    "    def compile(self, d_optimizer, g_optimizer, loss_fn):\n",
    "        super(GAN, self).compile()\n",
    "        self.d_optimizer = d_optimizer\n",
    "        self.g_optimizer = g_optimizer\n",
    "        self.loss_fn = loss_fn\n",
    "\n",
    "    def train_step(self, real_images):\n",
    "        if isinstance(real_images, tuple):\n",
    "            real_images = real_images[0]\n",
    "        # Sample random points in the latent space\n",
    "        batch_size = tf.shape(real_images)[0]\n",
    "        random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))\n",
    "\n",
    "        # Decode them to fake images\n",
    "        generated_images = self.generator(random_latent_vectors)\n",
    "\n",
    "        # Combine them with real images\n",
    "        combined_images = tf.concat([generated_images, real_images], axis=0)\n",
    "\n",
    "        # Assemble labels discriminating real from fake images\n",
    "        labels = tf.concat(\n",
    "            [tf.ones((batch_size, 1)), tf.zeros((batch_size, 1))], axis=0\n",
    "        )\n",
    "        # Add random noise to the labels - important trick!\n",
    "        labels += 0.05 * tf.random.uniform(tf.shape(labels))\n",
    "\n",
    "        # Train the discriminator\n",
    "        with tf.GradientTape() as tape:\n",
    "            predictions = self.discriminator(combined_images)\n",
    "            d_loss = self.loss_fn(labels, predictions)\n",
    "        grads = tape.gradient(d_loss, self.discriminator.trainable_weights)\n",
    "        self.d_optimizer.apply_gradients(\n",
    "            zip(grads, self.discriminator.trainable_weights)\n",
    "        )\n",
    "\n",
    "        # Sample random points in the latent space\n",
    "        random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))\n",
    "\n",
    "        # Assemble labels that say \"all real images\"\n",
    "        misleading_labels = tf.zeros((batch_size, 1))\n",
    "\n",
    "        # Train the generator (note that we should *not* update the weights\n",
    "        # of the discriminator)!\n",
    "        with tf.GradientTape() as tape:\n",
    "            predictions = self.discriminator(self.generator(random_latent_vectors))\n",
    "            g_loss = self.loss_fn(misleading_labels, predictions)\n",
    "        grads = tape.gradient(g_loss, self.generator.trainable_weights)\n",
    "        self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))\n",
    "        return {\"d_loss\": d_loss, \"g_loss\": g_loss}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ee4d8d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare the dataset. We use both the training & test MNIST digits.\n",
    "batch_size = 64\n",
    "(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()\n",
    "all_digits = np.concatenate([x_train, x_test])\n",
    "all_digits = all_digits.astype(\"float32\") / 255.0\n",
    "all_digits = np.reshape(all_digits, (-1, 28, 28, 1))\n",
    "dataset = tf.data.Dataset.from_tensor_slices(all_digits)\n",
    "dataset = dataset.shuffle(buffer_size=1024).batch(batch_size)\n",
    "\n",
    "gan = GAN(discriminator=discriminator, generator=generator, latent_dim=latent_dim)\n",
    "gan.compile(\n",
    "    d_optimizer=keras.optimizers.Adam(learning_rate=0.0003),\n",
    "    g_optimizer=keras.optimizers.Adam(learning_rate=0.0003),\n",
    "    loss_fn=keras.losses.BinaryCrossentropy(from_logits=True),\n",
    ")\n",
    "\n",
    "# To limit the execution time, we only train on 100 batches. You can train on\n",
    "# the entire dataset. You will need about 20 epochs to get nice results.\n",
    "gan.fit(dataset.take(100), epochs=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8726a905",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}