research/neural_clustering/examples/ncp_demo_with_mog_data.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S0g501NXWck7"
      },
      "source": [
        "##### Copyright 2021 Google LLC"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DkHNCmJ2Wjvu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZPZ7QUMPDTUc"
      },
      "source": [
        "# Learning to Cluster using Neural Clustering Process (NCP)\n",
        "\n",
        "Clustering has traditionally been posed as an unsupervised learning problem. However, unsupervised clustering methods often rely on strong assumptions about data geometry, and the clusters produced might not agree with human labels. Recently, supervised formulations of clustering have gained increasing interest, with novel objective functions and neural network architectures that directly learn complex clustering objectives from labeled data. In this tutorial, we will demonstrate a supervised neural clustering method, [Neural Clustering Process (NCP)](https://arxiv.org/pdf/1901.00409.pdf), which can be trained end-to-end to perform amortized probabilistic clustering without assuming a fixed or maximum number of clusters. Source code for the neural clustering models can be found on [Github](https://github.com/tensorflow/neural-structured-learning/tree/master/research/neural_clustering)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jr78lLlCLPIY"
      },
      "source": [
        "The high-level steps that we will follow in this colab are:\n",
        "1. Create synthetic data for clustering using a [Mixture of Gaussians (MOG)](https://en.wikipedia.org/wiki/Mixture_model#Gaussian_mixture_model) generative model.\n",
        "2. Define a NCP model architecture based on multi-layer perceptrons (MLP).\n",
        "3. Train the NCP model on synthetic MOG data.\n",
        "4. Perform probabilistic clustering using the trained NCP model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G8f4qAR4WQdk"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Install the `neural_clustering` module by running the code cell below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mUOEqgqNvGhI"
      },
      "outputs": [],
      "source": [
        "!pip install 'git+https://github.com/tensorflow/neural-structured-learning.git#egg=neural_clustering\u0026subdirectory=research/neural_clustering'\n",
        "!pip install tensorflow_addons"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WnkXf-51LieQ"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_addons as tfa\n",
        "from sklearn import metrics\n",
        "\n",
        "from neural_clustering import models\n",
        "from neural_clustering import data_generators"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "it6uioM7wh9U"
      },
      "outputs": [],
      "source": [
        "print(\"Num GPUs Available: \", len(tf.config.experimental.list_physical_devices('GPU')))\n",
        "COLORS = np.array(sns.color_palette() + sns.color_palette(\"pastel\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p1XCS8GVRfJF"
      },
      "source": [
        "## Generating and visualizing synthetic data\n",
        "\n",
        "The synthetic data will be generated in two steps. \n",
        "\n",
        "- First, we will produce random cluster partitions using a [Chinese Restaurant Process (CRP)](https://en.wikipedia.org/wiki/Chinese_restaurant_process) generative model. \n",
        "\n",
        "- Given the cluster partitions, synthetic data points in each cluster will then be generated using [Mixture of Gaussians (MOG)](https://en.wikipedia.org/wiki/Mixture_model#Gaussian_mixture_model). \n",
        "\n",
        "The code below instantiates a MOG data generator that produces 2 dimensional data so that we can visualize them easily. We also specify the discount parameter of CRP and the prior parameters of the mixture distribution."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cOeBcwekfw6y"
      },
      "outputs": [],
      "source": [
        "crp_generator = data_generators.CRPGenerator(alpha=0.8)\n",
        "mog_generator = data_generators.MOGGenerator(\n",
        "    partition_generator=crp_generator,\n",
        "    x_dim=2,\n",
        "    prior_sigma=10,\n",
        "    sigma_min=0.5,\n",
        "    sigma_max=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KcEpRSn1YZAW"
      },
      "source": [
        "Let's creates some synthetic data and plot them for visualization. \n",
        "\n",
        "In this example, we generate a mini-batch of four MOG samples. The parameter `N` specifies the number of data points in each sample. The plot shows that each sample contains a variable number of clusters and a different distribution of data points. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wFKmjlDbuLpR"
      },
      "outputs": [],
      "source": [
        "np.random.seed(8)\n",
        "N = 1000  #@param {\"type\": \"number\"}\n",
        "batch_size = 4\n",
        "\n",
        "mog_data, mog_labels = mog_generator.generate_batch(N, batch_size=batch_size)\n",
        "print(\"data shape:\", mog_data.shape)\n",
        "print(\"labels shape:\", mog_labels.shape)\n",
        "\n",
        "fig, axes = plt.subplots(1, batch_size, figsize=(batch_size * 4.5, 4))\n",
        "for i, ax in enumerate(axes):\n",
        "  ax.scatter(mog_data[i, :, 0], mog_data[i, :, 1], s=0.3, c=COLORS[mog_labels[i]])\n",
        "  ax.set_title(f'Sample {i}')\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DeGLzTtTHMWh"
      },
      "source": [
        "## Neural Clustering Process (NCP) model\n",
        "\n",
        "An end-to-end NCP clustering model contains the following two components:\n",
        "\n",
        "- **An encoder network** that transforms the input data into a fixed-dimensional vector. The type of encoder depends on the structure of the input data. For example, a convolutional neural network can be used for image data, and a graph neural network for graph data, etc. Here, we use a simple multi-layer perceptron for our MOG data.\n",
        "\n",
        "- **A NCP network** that takes as input the encoded vectors and learns to cluster these data points. The NCP algorithm assigns cluster labels to N data points sequentially by N forward passes throught the network. Our implementation decomposes the algorithm into two parts. The [`NCPBase`](https://github.com/tensorflow/neural-structured-learning/blob/master/research/neural_clustering/models/ncp_base.py) model class handles the forward pass for one iteration of NCP, and the [`NCPWrapper`](https://github.com/tensorflow/neural-structured-learning/blob/master/research/neural_clustering/models/ncp_wrapper.py) class wraps an [`NCPBase`](https://github.com/tensorflow/neural-structured-learning/blob/master/research/neural_clustering/models/ncp_base.py) model to perform the full sequential clustering algorithm. This is similar to the idea of [wrapping RNN cells into RNN layers](https://www.tensorflow.org/guide/keras/rnn). Both `NCPBase` and `NCPWrapper` are subclasses of `tf.keras.Model`.\n",
        "\n",
        "These two networks are then trained jointly through backpropagation.\n",
        "\n",
        "The [`NCPBase`](https://github.com/tensorflow/neural-structured-learning/blob/master/research/neural_clustering/models/ncp_base.py) model class itself does not define any neural network layer. Rather, it implements the NCP algorithm and expects the users to provide the neural network structures as its building-block layers. This provides maximal flexibility on the model architecture. In the code below, we use the [`NCPWithMLP`](https://github.com/tensorflow/neural-structured-learning/blob/master/research/neural_clustering/models/ncp_models.py) model, a subclass of [`NCPBase`](https://github.com/tensorflow/neural-structured-learning/blob/master/research/neural_clustering/models/ncp_base.py) that defines the building-block layers of NCP as multi-layer perceptrons (MLP)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0VqKY7zFzI_j"
      },
      "outputs": [],
      "source": [
        "# The default dimension size of hidden layers.\n",
        "layer_dim = 128\n",
        "\n",
        "# An encoder model for MOG data. Here we use a multi-layer perceptron (MLP).\n",
        "mog_encoder = tf.keras.Sequential([\n",
        "    tf.keras.layers.Dense(layer_dim//2, activation=\"relu\", name=\"e1\"),\n",
        "    tf.keras.layers.Dense(layer_dim//2, activation=\"relu\", name=\"e2\"),\n",
        "    tf.keras.layers.Dense(layer_dim, name=\"e3\"),\n",
        "])\n",
        "\n",
        "# A NCP model with multi-layer perceptrons (MLP) as its building blocks.\n",
        "ncp_model = models.NCPWithMLP(\n",
        "    assigned_point_layer_out_dim=layer_dim,\n",
        "    assigned_point_layer_hidden_dims=[layer_dim] * 2,\n",
        "    unassigned_point_layer_out_dim=layer_dim,\n",
        "    unassigned_point_layer_hidden_dims=[layer_dim] * 2,\n",
        "    cluster_layer_out_dim=layer_dim,\n",
        "    cluster_layer_hidden_dims=[layer_dim] * 2,\n",
        "    logits_layer_hidden_dims=[layer_dim] * 2)\n",
        "\n",
        "# The NCPWrapper class wraps the base model for the full sequential clustering\n",
        "# algorithm. The sampler parameter is used at inference time for predicting the\n",
        "# cluster labels. It's set to None during training because training labels are\n",
        "# available.\n",
        "ncp_wrapper_model = models.NCPWrapper(ncp_model=ncp_model, sampler=None)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "By5UjunozJ2A"
      },
      "source": [
        "## Training the NCP clustering model\n",
        "\n",
        "In order to train a NCP model given input data and target cluster labels, we minimize the negative log likelihood loss using the `logits` returned from `ncp_wrapper_model`. In the `train_step` function below, we define a custom training loop that jointly updates the weights of the encoder network and NCP network by gradient descent."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "viCXwWamzM--"
      },
      "outputs": [],
      "source": [
        "optimizer = tfa.optimizers.AdamW(learning_rate=0.001, weight_decay=0.0001)\n",
        "\n",
        "def train_step(inputs, targets):\n",
        "  with tf.GradientTape() as tape:\n",
        "    encoded = mog_encoder(inputs)\n",
        "    logits, _ = ncp_wrapper_model(encoded, targets, training=True)\n",
        "    loss = ncp_wrapper_model.loss_function(logits)\n",
        "\n",
        "  variables = ncp_wrapper_model.trainable_variables + mog_encoder.trainable_variables\n",
        "  gradients = tape.gradient(loss, variables)\n",
        "  optimizer.apply_gradients(zip(gradients, variables))\n",
        "  \n",
        "  # returns the loss averaged by the number of data points.\n",
        "  return loss / inputs.shape[1]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "STpGWHuYzOdr"
      },
      "source": [
        "Now let's train the clustering model. For the purpose of illustration, we use a small number of training steps.\n",
        "\n",
        "In each mini-batch, we generate synthetic training data similar to the examples shown earlier. The total number of data points to cluster in each training step is randomly sampled from a range, so that the model learns to cluster data of different sizes. The synthetic training examples across the mini-batch may also contain different number clusters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m0V_5a_hUKbW"
      },
      "outputs": [],
      "source": [
        "training_steps = 500  #@param {\"type\": \"number\"}\n",
        "n_points_min = 50  #@param {\"type\": \"number\"}\n",
        "n_points_max = 200  #@param {\"type\": \"number\"}\n",
        "batch_size = 32\n",
        "losses = []\n",
        "\n",
        "for i in range(training_steps):\n",
        "  n_total_points = np.random.randint(n_points_min, n_points_max)\n",
        "  batch_data, labels = mog_generator.generate_batch(n_total_points, batch_size=batch_size) \n",
        "\n",
        "  t_start = time.time()\n",
        "  loss = train_step(batch_data, labels)\n",
        "  batch_time = time.time() - t_start\n",
        "\n",
        "  losses.append(loss.numpy())\n",
        "  if i % 10 == 0:\n",
        "    print(f'Batch {i}, Number of data points: {n_total_points}, Loss: {np.mean(losses[-10:]):.3g}, Time/Batch: {batch_time:.2f}s')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FcjQFfY79cFR"
      },
      "outputs": [],
      "source": [
        "plt.plot(np.convolve(losses, np.ones(5)/5, mode='valid'))\n",
        "plt.ylabel(\"Loss\")\n",
        "plt.xlabel(\"Training Step\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eWW8TA-LeFY9"
      },
      "source": [
        "## Probabilistic clustering using trained NCP models\n",
        "\n",
        "The trained NCP model can be used to cluster new data points. Remember that NCP assigns cluster labels to each data point sequentially. In each iteration, NCP computes the unnormalized log probability (logits) of whether the new point should join one of the existing clusters or create a new one. This probability is then used for cluster assignment.\n",
        "\n",
        "Here, we demonstrate two methods for clustering assignment. The **greedy sampler** below always assigns a new point to the cluster with the highest probability, whereas the **random categorical sampler** will randomly sample a cluster label based on the categorical distribution defined by the logits. \n",
        "\n",
        "In the examples below, we will create a `ncp_sampler` by wrapping the trained `ncp_model` together with a sampler class, either the `GreedySampler` or the `CategoricalSampler`.\n",
        "\n",
        "We will generate some test data, and ask the trained NCP model to cluster them. The ground truth and inferred cluster labels will be visualized in plots. The joint probability of the final clustering configuration are computed using the logits from all cluster assignments, and shown in the plots below (e.g. `p=0.035`). We also evaluate the clustering quality using [adjusted mutual information (AMI)](https://en.wikipedia.org/wiki/Adjusted_mutual_information), which measures the agreement between the predicted cluster labels and ground truth. The AMI score ranges from 0 to 1, and a higher score indicates higher clustering quality."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b57lRa0Mfvla"
      },
      "outputs": [],
      "source": [
        "# Helper function for plotting the clustering results.\n",
        "def plot_clusters(data, labels, predictions, logits):\n",
        "  batch_size = data.shape[0]\n",
        "  joint_probs = np.exp(tf.reduce_sum(logits, axis=1))\n",
        "  ami_scores = [metrics.adjusted_mutual_info_score(labels[i], predictions[i]) for i in range(batch_size)]\n",
        "\n",
        "  _, axes = plt.subplots(2, batch_size, figsize=(4 * batch_size, 8))\n",
        "  for i in range(batch_size):\n",
        "    axes[0, i].scatter(data[i, :, 0], data[i, :, 1], s=3, c=COLORS[labels[i]])\n",
        "    axes[0, i].set_title(f'Test Example {i}\\nGround Truth')\n",
        "    axes[0, i].get_xaxis().set_visible(False)\n",
        "    axes[0, i].get_yaxis().set_visible(False)\n",
        "    axes[1, i].scatter(data[i, :, 0], data[i, :, 1], s=3, c=COLORS[predictions[i].numpy()])\n",
        "    axes[1, i].set_title(f'Inferred\\np = {joint_probs[i]:.4f}, AMI = {ami_scores[i]:.3f}')\n",
        "    axes[1, i].get_xaxis().set_visible(False)\n",
        "    axes[1, i].get_yaxis().set_visible(False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gOTgFU3Ef9Jc"
      },
      "source": [
        "### Test data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8ckBdWiNa-IP"
      },
      "outputs": [],
      "source": [
        "seed = 5  #@param {\"type\": \"number\"}\n",
        "np.random.seed(seed)\n",
        "n_points = 200  #@param {\"type\": \"number\"}\n",
        "batch_size = 4\n",
        "test_data, test_labels = mog_generator.generate_batch(n_points, batch_size=batch_size) "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jTqrZ6YXHWg6"
      },
      "source": [
        "### Clustering using NCP + greedy sampler\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rd7lPHFVU4m-"
      },
      "outputs": [],
      "source": [
        "ncp_greedy_sampler = models.NCPWrapper(ncp_model=ncp_model, sampler=models.GreedySampler())\n",
        "\n",
        "test_encoded = mog_encoder(test_data)\n",
        "logits, test_pred = ncp_greedy_sampler(test_encoded, training=False)\n",
        "\n",
        "plot_clusters(test_data, test_labels, test_pred, logits)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SOsKLx0RHfAd"
      },
      "source": [
        "### Clustering using NCP + random categorical sampler"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3nYQU2wMihZy"
      },
      "outputs": [],
      "source": [
        "ncp_categorical_sampler = models.NCPWrapper(ncp_model=ncp_model, sampler=models.CategoricalSampler())\n",
        "\n",
        "test_encoded = mog_encoder(test_data)\n",
        "logits, test_pred = ncp_categorical_sampler(test_encoded, training=False)\n",
        "\n",
        "plot_clusters(test_data, test_labels, test_pred, logits)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IWz7WhCblk6j"
      },
      "source": [
        "### Capturing clustering uncertainty by parallelized sampling\n",
        "\n",
        "In a single run, the greedy sampler usually produces clustering configurations with higher probabilities and AMI scores than the random categorical sampler. However, the random categorical sampler can be used to sample many sets of cluster labels in parallel, and by choosing the output with the highest probability, we may get higher quality clusters than by using the greedy sampler. \n",
        "\n",
        "Importantly, this random sampling capability of NCP can be used to capture the uncertainty of cluster assignment by generating multiple plausible clustering configurations.\n",
        "\n",
        "In the example below, we create one test example, and ask NCP to sample multiple sets of cluster labels. The top few clustering results with high probability are shown in the plots."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iLM884vYqHlR"
      },
      "outputs": [],
      "source": [
        "# Test data.\n",
        "seed = 1  #@param {\"type\": \"number\"}\n",
        "np.random.seed(seed)\n",
        "n_points = 200  #@param {\"type\": \"number\"}\n",
        "test_data, test_labels = mog_generator.generate_batch(n_points, batch_size=1) \n",
        "\n",
        "test_encoded = mog_encoder(test_data)\n",
        "\n",
        "# The encoded tensor is replicated to perform parallelized sampling of cluster labels.\n",
        "n_parallel_samples = 100  #@param {\"type\": \"number\"}\n",
        "test_encoded_replicated = tf.repeat(test_encoded, repeats=n_parallel_samples, axis=0)\n",
        "\n",
        "logits, test_pred = ncp_categorical_sampler(test_encoded_replicated, training=False)\n",
        "joint_probs = np.exp(tf.reduce_sum(logits, axis=1))\n",
        "\n",
        "# Sorting the unique predictions by probability.\n",
        "# (np.unique returns sorted unique elements.)\n",
        "joint_probs, unique_index = np.unique(joint_probs, return_index=True)\n",
        "joint_probs, unique_index = joint_probs[::-1], unique_index[::-1]\n",
        "test_pred = tf.gather(test_pred, unique_index)\n",
        "\n",
        "# Only plotting the top few predictions.\n",
        "top_n = min(4, len(joint_probs))\n",
        "\n",
        "ami_scores = [metrics.adjusted_mutual_info_score(test_labels[0], test_pred[i]) for i in range(top_n)]\n",
        "\n",
        "_, axes = plt.subplots(1, top_n + 1, figsize=(4 * (top_n + 1), 4))\n",
        "axes[0].scatter(test_data[0, :, 0], test_data[0, :, 1], s=3, c=COLORS[test_labels[0]])\n",
        "axes[0].set_title(f'Ground Truth')\n",
        "\n",
        "for i in range(top_n):\n",
        "  axes[i+1].scatter(test_data[0, :, 0], test_data[0, :, 1], s=3, c=COLORS[test_pred[i].numpy()])\n",
        "  axes[i+1].set_title(f'Inferred #{i+1}\\np = {joint_probs[i]:.4f}, AMI = {ami_scores[i]:.3f}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zkJuAW76qdh_"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "In this colab, we have created a Neural Clustering Process (NCP) model, and trained the model to cluster synthetic data. We have seen that NCP can learn to perform probabilistic clustering on datasets of different sizes and containing different number of clusters."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "Neural Clustering Process demo with synthetic data",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}