shape_segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Shape segmentation\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 segmentation 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 (HarmonicConv, HarmonicResNetBlock,\n",
    "                ParallelTransportPool, ParallelTransportUnpool,\n",
    "                ComplexLin, ComplexNonLin)\n",
    "# Utility functions\n",
    "from utils.harmonic import magnitudes\n",
    "# Rotated MNIST dataset\n",
    "from datasets import ShapeSeg\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\n",
    "\n",
    "# Number of classes for segmentation\n",
    "n_classes = 8"
   ]
  },
  {
   "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', 'ShapeSeg')\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', sample_n=1024),\n",
    "    VectorHeat(),\n",
    "    Subsample(),\n",
    "    NormalizeAxes()\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 = ShapeSeg(path, False, pre_transform=pre_transform)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size)\n",
    "train_dataset = ShapeSeg(path, True, pre_transform=pre_transform, transform=transform)\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)"
   ]
  },
  {
   "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/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",
    "        # Stack 1\n",
    "        self.resnet_block11 = HarmonicResNetBlock(nf[0], nf[0], max_order, n_rings, prev_order=0)\n",
    "        self.resnet_block12 = HarmonicResNetBlock(nf[0], nf[0], max_order, n_rings)\n",
    "\n",
    "        # Pool\n",
    "        self.pool = ParallelTransportPool(1, scale1_transform)\n",
    "\n",
    "        # Stack 2\n",
    "        self.resnet_block21 = HarmonicResNetBlock(nf[0], nf[1], max_order, n_rings)\n",
    "        self.resnet_block22 = HarmonicResNetBlock(nf[1], nf[1], max_order, n_rings)\n",
    "\n",
    "        # Stack 3\n",
    "        self.resnet_block31 = HarmonicResNetBlock(nf[1], nf[1], max_order, n_rings)\n",
    "        self.resnet_block32 = HarmonicResNetBlock(nf[1], nf[1], max_order, n_rings)\n",
    "\n",
    "        # Unpool\n",
    "        self.unpool = ParallelTransportUnpool(from_lvl=1)\n",
    "\n",
    "        # Stack 4\n",
    "        self.resnet_block41 = HarmonicResNetBlock(nf[1] + nf[0], nf[0], max_order, n_rings)\n",
    "        self.resnet_block42 = HarmonicResNetBlock(nf[0], nf[0], 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[0], 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 positions to nf[0] features\n",
    "        x = F.relu(self.lin0(x))\n",
    "        \n",
    "        # Convert input features into complex numbers\n",
    "        x = torch.stack((x, torch.zeros_like(x)), dim=-1).unsqueeze(1)\n",
    "\n",
    "        # Stack 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",
    "        x_prepool = self.resnet_block12(x, *attributes)\n",
    "\n",
    "        # Pooling\n",
    "        # Apply parallel transport pooling\n",
    "        x, data, data_pooled = self.pool(x_prepool, data)\n",
    "        \n",
    "        # Stack 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",
    "        x = self.resnet_block22(x, *attributes_pooled)\n",
    "\n",
    "        # Stack 3\n",
    "        x = self.resnet_block31(x, *attributes_pooled)\n",
    "        x = self.resnet_block32(x, *attributes_pooled)\n",
    "\n",
    "        # Unpooling\n",
    "        x = self.unpool(x, data)\n",
    "        # Concatenate pre-pooling x with post-pooling x\n",
    "        x = torch.cat((x, x_prepool), dim=2)\n",
    "\n",
    "        # Stack 3\n",
    "        x = self.resnet_block41(x, *attributes)\n",
    "        x = self.resnet_block42(x, *attributes)\n",
    "\n",
    "        x = self.conv_final(x, *attributes)\n",
    "\n",
    "        # Take radial component from features and sum streams\n",
    "        x = magnitudes(x, keepdim=False)\n",
    "        x = x.sum(dim=1)\n",
    "\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.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",
    "    if epoch > 20:\n",
    "        for param_group in optimizer.param_groups:\n",
    "            param_group['lr'] = 0.001\n",
    "\n",
    "    for data in progressbar.progressbar(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 += data.y.size(0)\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, may take a while...')\n",
    "# Try with fewer epochs if you're in a timecrunch\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
}