04-spiral_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spiral classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "from math import pi as π"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from res.plot_lib import plot_data, plot_model, plot_embeddings, set_default, acc, overwrite, plot_2d_energy_levels, plot_3d_energy_levels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_default()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = 12345\n",
    "torch.manual_seed(seed)\n",
    "N = 1000  # num_samples_per_class\n",
    "n = 2     # input dimensions\n",
    "K = 5     # num_classes\n",
    "d = 100   # num_hidden_units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate spirals\n",
    "\n",
    "t = torch.linspace(0, 1, N)\n",
    "a = 0.8 * t + 0.2  # amplitude 0.2 → 1.0\n",
    "X = list()\n",
    "y = list()\n",
    "for k in range(K):\n",
    "    θ = (2 * t + k) * 2 * π / K + 0.2 * torch.randn(N)\n",
    "    X.append(torch.stack((a * θ.sin(), a * θ.cos()), dim=1))\n",
    "    y.append(torch.zeros(N, dtype=torch.long).fill_(k))\n",
    "X = torch.cat(X)\n",
    "y = torch.cat(y)\n",
    "\n",
    "print(\"Shapes:\")\n",
    "print(\"X:\", tuple(X.size()))\n",
    "print(\"y:\", tuple(y.size()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# And visualise them\n",
    "plot_data(X, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build and train a neural net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 1e-3\n",
    "lambda_l2 = 1e-5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model definition\n",
    "model = nn.Sequential(\n",
    "    nn.Linear(n, d),\n",
    "#     nn.ReLU(),  # Comment this line for a linear model\n",
    "    nn.Linear(d, K)  # (Optional) Comment this line and uncomment the next one to display 2D embeddings below\n",
    "#     nn.Linear(d, 2), nn.Linear(2, K)\n",
    ")\n",
    "model.to(device)  # possibly send to CUDA\n",
    "\n",
    "# Cross entropy given the output logits\n",
    "C = torch.nn.CrossEntropyLoss(reduction='none')\n",
    "\n",
    "# Using Adam optimiser\n",
    "optimiser = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=lambda_l2) # built-in L2\n",
    "\n",
    "# Full-batch training loop\n",
    "for t in range(2_000):\n",
    "    \n",
    "    # Feed forward to get the logits\n",
    "    l = model(X)\n",
    "    \n",
    "    # Compute the free energy F\n",
    "    F = C(l, y)\n",
    "    L = F.mean()\n",
    "    \n",
    "    # Zero the gradients\n",
    "    optimiser.zero_grad()\n",
    "    \n",
    "    # Backward pass to compute and accumulate the gradient\n",
    "    # of the free energy w.r.t our learnable params\n",
    "    L.backward()\n",
    "    \n",
    "    # Update params\n",
    "    optimiser.step()\n",
    "    \n",
    "    # Display epoch, L, and accuracy\n",
    "    overwrite(f'[EPOCH]: {t}, [LOSS]: {L.item():.6f}, [ACCURACY]: {acc(l, y):.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot trained model\n",
    "print(model)\n",
    "plot_model(X, y, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# (Optional) Plot internal 2D embeddings if available\n",
    "plot_embeddings(X, y, model, zoom=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute logits for a fine grid over the input space\n",
    "\n",
    "mesh = torch.arange(-1.5, 1.5, 0.01)\n",
    "xx, yy = torch.meshgrid(mesh, mesh)\n",
    "grid = torch.stack((xx.reshape(-1), yy.reshape(-1)), dim=1)\n",
    "with torch.no_grad():\n",
    "    logits = model(grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choice of free energy\n",
    "\n",
    "fe = 'cross-entropy'\n",
    "# fe = 'negative logit'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Switch to non-interactive matplotlib\n",
    "%matplotlib inline\n",
    "set_default()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! mkdir {m}-levels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Plot 2d energy levels\n",
    "\n",
    "for k in range(K):\n",
    "    if fe == 'cross-entropy':\n",
    "        F = C(logits, torch.LongTensor(1).fill_(k).expand(logits.size(0)))\n",
    "        F = F.reshape(xx.shape)\n",
    "        plot_2d_energy_levels(X, y, (xx, yy, F, k, K), (0, 35), (1, 35, 4))\n",
    "\n",
    "    elif fe == 'negative logit':\n",
    "        F = -logits[:, k]\n",
    "        F = F.reshape(xx.shape)\n",
    "        plot_2d_energy_levels(X, y, (xx, yy, F, k, K), (-20, 20), (-20, 21, 2.5))\n",
    "        \n",
    "#     plt.savefig(f'{m}-levels/{k}.png', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! ffmpeg -framerate 1 -i {m}-levels/%d.png -r 25 -vf \"crop=trunc(iw/2)*2:trunc(ih/2)*2\" -pix_fmt yuv420p {m}-levels.mp4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Switch to interactive matplotlib\n",
    "%matplotlib notebook\n",
    "set_default()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cross-entropy\n",
    "if fe == 'cross-entropy':\n",
    "    fig, ax = plot_3d_energy_levels(X, y, (xx, yy, F, k, K), (0, 18), (0, 19, 1), (0, 19, 2))\n",
    "elif fe == 'negative logit':\n",
    "    fig, ax = plot_3d_energy_levels(X, y, (xx, yy, F, k, K), (-30, 20), (-30, 20, 1), (-30, 21, 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! mkdir {m}-3d-levels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Spin it around (and maybe save to disk)\n",
    "δ = 10\n",
    "for angle in range(0, 360, δ):\n",
    "    ax.view_init(30, -60 + angle)\n",
    "    fig.canvas.draw()\n",
    "#     plt.pause(.001)\n",
    "#     plt.savefig(f'{m}-3d-levels/{angle:03d}.png', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ! ffmpeg -i {m}-3d-levels/%03d.png -vf \"crop=trunc(iw/2)*2:trunc(ih/2)*2\" -pix_fmt yuv420p {m}-3d-levels.mp4"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pDL] *",
   "language": "python",
   "name": "conda-env-pDL-py"
  },
  "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.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}