rotated_mnist.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rotated MNIST\n",
    "\n",
    "The notebooks in this folder replicate the experiments as performed for [CNNs on Surfaces using Rotation-Equivariant Features](https://doi.org/10.1145/3386569.3392437).\n",
    "\n",
    "The current notebook replicates the Rotated MNIST experiments from section `5.3 Evaluation`.\n",
    "\n",
    "## Imports\n",
    "We start by importing dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# File reading and progressbar\n",
    "import os.path as osp\n",
    "import progressbar\n",
    "\n",
    "# PyTorch and PyTorch Geometric dependencies\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch_geometric.transforms as T\n",
    "from torch_geometric.nn import global_mean_pool\n",
    "from torch_geometric.nn.inits import zeros\n",
    "from torch_geometric.data import DataLoader\n",
    "\n",
    "# Harmonic Surface Networks components\n",
    "# Layers\n",
    "from nn import HarmonicConv, ComplexNonLin, ParallelTransportPool\n",
    "# Utility functions\n",
    "from utils.harmonic import magnitudes, c_batch_norm\n",
    "# Rotated MNIST dataset\n",
    "from datasets import MNISTSphere\n",
    "# Transforms\n",
    "from transforms import VectorHeat, HarmonicPrecomp, MultiscaleRadiusGraph, ScaleMask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings\n",
    "Next, we set a few parameters for our network. You can change these settings to experiment with different configurations of the network. Right now, the settings are set to the ones used in the paper."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Maximum rotation order for streams\n",
    "max_order = 1\n",
    "\n",
    "# Number of rings in the radial profile\n",
    "n_rings = 4\n",
    "\n",
    "# Number of filters per block\n",
    "nf = [8, 16, 32]\n",
    "\n",
    "# Ratios used for pooling\n",
    "ratios=[1, 0.5, 0.25]\n",
    "\n",
    "# Radius of convolution for each scale\n",
    "radii = [0.3, 0.45, 0.8]\n",
    "\n",
    "# Number of datasets per batch\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset\n",
    "To get our dataset ready for training, we need to perform the following steps:\n",
    "1. Provide a path to load and store the dataset.\n",
    "2. Define transformations to be performed on the dataset:\n",
    "    - A transformation that computes a multi-scale radius graph and precomputes the logarithmic map.\n",
    "    - A transformation that masks the edges and vertices per scale and precomputes convolution components.\n",
    "3. Assign and load the datasets.\n",
    "\n",
    "Note that for the sphere, the logarithmic map and precomputed components are equal for every datapoint. Hence, we cache these computations and reuse them for every shape. To store the caches, we create a cache folder first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Provide a path to load and store the dataset.\n",
    "# Make sure that you have created a folder 'data' somewhere\n",
    "# and that you have downloaded and moved the raw dataset there\n",
    "path = osp.join('data', 'MNISTSphere')\n",
    "\n",
    "# 2. Define transformations to be performed on the dataset:\n",
    "# Transformation that computes a multi-scale radius graph and precomputes the logarithmic map.\n",
    "transform = T.Compose((\n",
    "    MultiscaleRadiusGraph(ratios, radii, loop=True, flow='target_to_source', cache_file='cache/radius_sphere.pt'),\n",
    "    VectorHeat(cache_file='cache/logmap_sphere.pt')\n",
    "))\n",
    "# Transformations that mask the edges and vertices per scale and precomputes convolution components.\n",
    "scale0_transform = T.Compose((\n",
    "    ScaleMask(0),\n",
    "    HarmonicPrecomp(n_rings, max_order, max_r=radii[0], cache_file='cache/p0_sphere.pt')\n",
    "))\n",
    "scale1_transform = T.Compose((\n",
    "    ScaleMask(1),\n",
    "    HarmonicPrecomp(n_rings, max_order, max_r=radii[1], cache_file='cache/p1_sphere.pt')\n",
    "))\n",
    "scale2_transform = T.Compose((\n",
    "    ScaleMask(2),\n",
    "    HarmonicPrecomp(n_rings, max_order, max_r=radii[2], cache_file='cache/p2_sphere.pt')\n",
    "))\n",
    "\n",
    "# 3. Assign and load the datasets.\n",
    "train_dataset = MNISTSphere(path, True, transform=transform)\n",
    "test_dataset = MNISTSphere(path, False, transform=transform)\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Network architecture\n",
    "Now, we create the network architecture by creating a new `nn.Module`, `Net`. We first setup each layer in the `__init__` method of the `Net` class and define the steps to perform for each batch in the `forward` method. The following figure shows a schematic of the architecture we will be implementing:\n",
    "\n",
    "<img src=\"img/classification_architecture.png\" width=\"500px\" />\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        # Block 1, scale 0\n",
    "        # Because we start with an m=0 input, we set prev_order to 0\n",
    "        self.conv1 = HarmonicConv(1, nf[0], max_order, n_rings, prev_order=0)\n",
    "        self.nonlin1 = ComplexNonLin(nf[0])\n",
    "        self.conv2 = HarmonicConv(nf[0], nf[0], max_order, n_rings)\n",
    "        self.bn1 = nn.BatchNorm1d((max_order + 1) * nf[0], eps=1e-3, momentum=0.01)\n",
    "\n",
    "        # Pooling to scale 1\n",
    "        self.pool1 = ParallelTransportPool(1, scale1_transform)\n",
    "\n",
    "        # Block 2, scale 1\n",
    "        self.conv3 = HarmonicConv(nf[0], nf[1], max_order, n_rings)\n",
    "        self.nonlin3 = ComplexNonLin(nf[1])\n",
    "        self.conv4 = HarmonicConv(nf[1], nf[1], max_order, n_rings)\n",
    "        self.bn2 = nn.BatchNorm1d((max_order + 1) * nf[1], eps=1e-3, momentum=0.01)\n",
    "\n",
    "        # Pooling to scale 2\n",
    "        self.pool2 = ParallelTransportPool(2, scale2_transform)\n",
    "\n",
    "        # Block 3, scale 2\n",
    "        self.conv5 = HarmonicConv(nf[1], nf[2], max_order, n_rings)\n",
    "        self.nonlin5 = ComplexNonLin(nf[2])\n",
    "        self.conv6 = HarmonicConv(nf[2], nf[2], max_order, n_rings)\n",
    "        self.bn3 = nn.BatchNorm1d((max_order + 1) * nf[2], eps=1e-3, momentum=0.01)\n",
    "\n",
    "        # Final Harmonic Convolution\n",
    "        # We set offset to False, \n",
    "        # because we will only use the radial component of the features after this\n",
    "        self.conv7 = HarmonicConv(nf[2], 10, max_order, n_rings, offset=False)\n",
    "        self.bias = nn.Parameter(torch.Tensor(10))\n",
    "        zeros(self.bias)\n",
    "\n",
    "    def forward(self, data):\n",
    "        # The input x is fed to our convolutional layers as a complex number and organized by rotation orders.\n",
    "        # Resulting matrix: [batch_size, max_order + 1, channels, complex]\n",
    "        x = torch.stack((data.x, torch.zeros_like(data.x)), dim=-1).unsqueeze(1)\n",
    "        batch_size = data.num_graphs\n",
    "        n_nodes = x.size(0)\n",
    "\n",
    "        # Block 1, scale 0\n",
    "        # Mask correct edges and nodes\n",
    "        data_scale0 = scale0_transform(data)\n",
    "        # Get edge indices and precomputations for scale 0\n",
    "        attributes = (data_scale0.edge_index, data_scale0.precomp, data_scale0.connection)\n",
    "        # Apply convolutions\n",
    "        x = self.conv1(x, attributes[0], attributes[1])\n",
    "        x = self.nonlin1(x)\n",
    "        x = self.conv2(x, *attributes)\n",
    "        x = c_batch_norm(x, batch_size, self.bn1, F.relu)\n",
    "        \n",
    "        # Pooling to scale 1\n",
    "        x, data, data_pooled = self.pool1(x, data)\n",
    "        # Get edge indices and precomputations for scale 1\n",
    "        attributes_pooled = (data_pooled.edge_index, data_pooled.precomp, data_pooled.connection)\n",
    "\n",
    "        # Block 2, scale 1\n",
    "        x = self.conv3(x, *attributes_pooled)\n",
    "        x = self.nonlin3(x)\n",
    "        x = self.conv4(x, *attributes_pooled)\n",
    "        x = c_batch_norm(x, batch_size, self.bn2, F.relu)\n",
    "\n",
    "        # Pooling to scale 2\n",
    "        x, data, data_pooled = self.pool2(x, data)\n",
    "        # Get edge indices and precomputations for scale 2\n",
    "        attributes_pooled = (data_pooled.edge_index, data_pooled.precomp, data_pooled.connection)\n",
    "\n",
    "        # Block 3, scale 2\n",
    "        x = self.conv5(x, *attributes_pooled)\n",
    "        x = self.nonlin5(x)\n",
    "        x = self.conv6(x, *attributes_pooled)\n",
    "        x = c_batch_norm(x, batch_size, self.bn3, F.relu)\n",
    "\n",
    "        # Final convolution\n",
    "        x = self.conv7(x, *attributes_pooled)\n",
    "        # Take radial component of each complex feature\n",
    "        x = magnitudes(x, keepdim=False)\n",
    "        # Sum the two streams\n",
    "        x = x.sum(dim=1)\n",
    "\n",
    "        # Global mean pool to retrieve classification\n",
    "        x = global_mean_pool(x, data.batch)\n",
    "        x = x + self.bias\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "Phew, we're through the hard part. Now, let's get to training. First, move the network to the GPU and setup an optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We want to train on a GPU. It'll take a long time on a CPU\n",
    "device = torch.device('cuda')\n",
    "# Move the network to the GPU\n",
    "model = Net().to(device)\n",
    "# Set up the ADAM optimizer with learning rate of 0.0076 (as used in H-Nets)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0076)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, define a training and test function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def train(epoch):\n",
    "    # Set model to 'train' mode\n",
    "    model.train()\n",
    "\n",
    "    for param_group in optimizer.param_groups:\n",
    "        # Slowly decrease the learning rate every epoch\n",
    "        param_group['lr'] = param_group['lr'] * np.power(0.1, epoch / 50)\n",
    "\n",
    "    for data in train_loader:\n",
    "        # Move training data to the GPU and optimize parameters\n",
    "        optimizer.zero_grad()\n",
    "        F.nll_loss(model(data.to(device)), data.y).backward()\n",
    "        optimizer.step()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train for 100 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('This may take a while...')\n",
    "# Try with fewer epochs if you're in a timecrunch\n",
    "for epoch in progressbar.progressbar(range(100), redirect_stdout=True):\n",
    "    train(epoch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing\n",
    "Finally, we test our model on the test dataset. Setup a test function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test():\n",
    "    # Set model to 'evaluation' mode\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "\n",
    "    for data in progressbar.progressbar(test_loader):\n",
    "        # Move test data to the GPU and return a prediction\n",
    "        data = data.to(device)\n",
    "        pred = model(data).max(1)[1]\n",
    "        correct += pred.eq(data.y).sum().item()\n",
    "    # Return the fraction of correctly classified shapes\n",
    "    return correct / len(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And return the accuracy on the test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_acc = test()\n",
    "print('Test: {:.6f}'.format(test_acc))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}