shape_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Shape classification\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 shape classification experiments from section `5.2 Comparisons`.\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.data import DataLoader\n",
    "from torch_geometric.nn.inits import zeros\n",
    "\n",
    "# Harmonic Surface Networks components\n",
    "# Layers\n",
    "from nn import (HarmonicResNetBlock, HarmonicConv,\n",
    "                ParallelTransportPool, ParallelTransportUnpool,\n",
    "                ComplexNonLin)\n",
    "# Utility functions\n",
    "from utils.harmonic import magnitudes\n",
    "# Rotated MNIST dataset\n",
    "from datasets import Shrec16\n",
    "# Transforms\n",
    "from transforms import (HarmonicPrecomp, VectorHeat, MultiscaleRadiusGraph, \n",
    "                        ScaleMask, FilterNeighbours, NormalizeArea, NormalizeAxes, Subsample)"
   ]
  },
  {
   "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 = 6\n",
    "\n",
    "# Number of filters per block\n",
    "nf = [16, 32]\n",
    "\n",
    "# Ratios used for pooling\n",
    "ratios=[1, 0.25]\n",
    "\n",
    "# Radius of convolution for each scale\n",
    "radii = [0.2, 0.4]\n",
    "\n",
    "# Number of datasets per batch\n",
    "batch_size = 1"
   ]
  },
  {
   "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."
   ]
  },
  {
   "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 datasets there\n",
    "path = osp.join('data', 'shrec')\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",
    "pre_transform = T.Compose((\n",
    "    NormalizeArea(),\n",
    "    MultiscaleRadiusGraph(ratios, radii, loop=True, flow='target_to_source'),\n",
    "    VectorHeat()\n",
    "))\n",
    "# Apply a random scale and random rotation to each shape\n",
    "transform = T.Compose((\n",
    "    T.RandomScale((0.85, 1.15)),\n",
    "    T.RandomRotate(45, axis=0),\n",
    "    T.RandomRotate(45, axis=1),\n",
    "    T.RandomRotate(45, axis=2))\n",
    ")\n",
    "\n",
    "# Transformations that masks the edges and vertices per scale and precomputes convolution components.\n",
    "scale0_transform = T.Compose((\n",
    "    ScaleMask(0),\n",
    "    FilterNeighbours(radii[0]),\n",
    "    HarmonicPrecomp(n_rings, max_order, max_r=radii[0]))\n",
    ")\n",
    "scale1_transform = T.Compose((\n",
    "    ScaleMask(1),\n",
    "    FilterNeighbours(radii[1]),\n",
    "    HarmonicPrecomp(n_rings, max_order, max_r=radii[1]))\n",
    ")\n",
    "\n",
    "# 3. Assign and load the datasets.\n",
    "test_dataset = Shrec16(path, False, pre_transform=pre_transform)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size)\n",
    "train_dataset = Shrec16(path, True, pre_transform=pre_transform, transform=transform)\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "\n",
    "n_classes = train_dataset.num_classes"
   ]
  },
  {
   "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. We use the left part of the U-ResNet architecture, with half of the ResNet blocks from the following figure:\n",
    "\n",
    "<img src=\"img/resnet_architecture.png\" width=\"800px\" />\n",
    "\n",
    "Let's get started!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        self.lin0 = nn.Linear(3, nf[0])\n",
    "\n",
    "        # ResNet Block 1\n",
    "        self.resnet_block11 = HarmonicResNetBlock(nf[0], nf[0], max_order, n_rings, prev_order=0)\n",
    "\n",
    "        # Pool\n",
    "        self.pool = ParallelTransportPool(1, scale1_transform)\n",
    "\n",
    "        # ResNet Block 2\n",
    "        self.resnet_block21 = HarmonicResNetBlock(nf[0], nf[1], max_order, n_rings)\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.conv_final = HarmonicConv(nf[1], n_classes, max_order, n_rings, offset=False)\n",
    "        \n",
    "        self.bias = nn.Parameter(torch.Tensor(n_classes))\n",
    "        zeros(self.bias)\n",
    "\n",
    "\n",
    "    def forward(self, data):\n",
    "        x = data.pos\n",
    "        \n",
    "        # Linear transformation from input descriptors to nf[0] features\n",
    "        x = F.relu(self.lin0(x))\n",
    "        # Convert input features into complex numbers\n",
    "        x = torch.stack((x, torch.zeros_like(x)), dim=-1).unsqueeze(1)\n",
    "\n",
    "        # ResNet Block 1\n",
    "        # Select only the edges and precomputed components of the first scale\n",
    "        data_scale0 = scale0_transform(data)\n",
    "        attributes = (data_scale0.edge_index, data_scale0.precomp, data_scale0.connection)\n",
    "        x = self.resnet_block11(x, *attributes)\n",
    "\n",
    "        # Pooling\n",
    "        # Apply parallel transport pooling\n",
    "        x, data, data_pooled = self.pool(x, data)\n",
    "        \n",
    "        # ResNet Block 2\n",
    "        # Store edge_index and precomputed components of the second scale\n",
    "        attributes_pooled = (data_pooled.edge_index, data_pooled.precomp, data_pooled.connection)\n",
    "        x = self.resnet_block21(x, *attributes_pooled)\n",
    "\n",
    "        x = self.conv_final(x, *attributes_pooled)\n",
    "        \n",
    "        # Take radial component from features and sum streams\n",
    "        x = magnitudes(x, keepdim=False)\n",
    "        x = x.sum(dim=1)\n",
    "\n",
    "        # Global mean pool\n",
    "        x = torch.mean(x, dim=0, keepdim=True)\n",
    "        x = x + self.bias\n",
    "        return F.log_softmax(x, dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to test the network without applying parallel transport, remove the `connection` argument from convolutions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training and testing\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.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, define a training and test function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch):\n",
    "    # Set model to 'train' mode\n",
    "    model.train()\n",
    "\n",
    "    # For the two stream architecture,\n",
    "    # it helps to decrease the learning rate earlier to stabilise training\n",
    "    if epoch > 30:\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group['lr'] = 0.001\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()\n",
    "        \n",
    "def test():\n",
    "    # Set model to 'evaluation' mode\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total_num = 0\n",
    "    for i, data in enumerate(test_loader):\n",
    "        pred = model(data.to(device)).max(1)[1]\n",
    "        correct += pred.eq(data.y).sum().item()\n",
    "        total_num += 1\n",
    "    return correct / total_num"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Train for 50 epochs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Start training...')\n",
    "for epoch in range(50):\n",
    "    train(epoch)\n",
    "    test_acc = test()\n",
    "    print(\"Epoch {} - Test: {:06.4f}\".format(epoch, 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
}