segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "label2name = {\n",
    "    0: 'Airplane   ',\n",
    "    1: 'Bag        ',\n",
    "    2: 'Cap        ',\n",
    "    3: 'Car        ',\n",
    "    4: 'Chair      ',\n",
    "    5: 'Earphone   ',\n",
    "    6: 'Guitar     ',\n",
    "    7: 'Knife      ',\n",
    "    8: 'Lamp       ',\n",
    "    9: 'Laptop     ',\n",
    "    10: 'Motorbike ',\n",
    "    11: 'Mug       ',\n",
    "    12: 'Pistol    ',\n",
    "    13: 'Rocket    ',\n",
    "    14: 'Skateboard',\n",
    "    15: 'Table     '\n",
    "}\n",
    "\n",
    "config = {\n",
    "    # General\n",
    "    'batchsize': 64,\n",
    "    'shuffle': True,\n",
    "    # Augmentations\n",
    "    'flip': False,\n",
    "    'ascale': True, 'as_min': 0.6667, 'as_max': 1.5,\n",
    "    'rotate': False, 'r_positions': 12, 'test_pos': None,\n",
    "    'translate': True, 't_rate': 0.1,\n",
    "    # Point clouds and kd-trees generation\n",
    "    'steps': 12, # also control the depth of the network\n",
    "    'dim': 3,\n",
    "    'lim': 1,\n",
    "    'det': False,\n",
    "    'gamma': 10.,\n",
    "    # NN options\n",
    "    'input_features': 'all', # 'all' for point coordinates, 'no' for feeding 1's as point features\n",
    "    'n_f': [64, \n",
    "            64,  64,  64,  64,  \n",
    "            128, 128, 128, 128, \n",
    "            256, 256, 256, 256, \n",
    "            512], # representation sizes\n",
    "    'n_mlp': [128, 128], # mlp sizes\n",
    "    'n_output': 50,\n",
    "    'l2': 1e-4,\n",
    "    'lr': 1e-4,\n",
    "    'n_ens': 16\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import h5py as h5\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "with h5.File('./shapenet_partanno.h5', 'r') as fin:\n",
    "    train_labels = np.array(fin['train_labels'])\n",
    "    train_borders = np.array(fin['train_borders'])\n",
    "    train_points = np.array(fin['train_points'])\n",
    "    train_point_labels = np.array(fin['train_point_labels'])\n",
    "    \n",
    "    val_labels = np.array(fin['val_labels'])\n",
    "    val_borders = np.array(fin['val_borders'])\n",
    "    val_points = np.array(fin['val_points'])\n",
    "    val_point_labels = np.array(fin['val_point_labels'])\n",
    "\n",
    "    test_labels = np.array(fin['test_labels'])\n",
    "    test_borders = np.array(fin['test_borders'])\n",
    "    test_points = np.array(fin['test_points'])\n",
    "    test_point_labels = np.array(fin['test_point_labels'])\n",
    "    \n",
    "label2point_labels = {}\n",
    "for i in xrange(train_labels.max()+1):\n",
    "    ind_labels = np.empty((0,), dtype=np.int8)\n",
    "    for ind in (train_labels == i).nonzero()[0]:\n",
    "        ind_labels = np.hstack((ind_labels, train_point_labels[train_borders[ind]:train_borders[ind+1]]))\n",
    "    for ind in (val_labels == i).nonzero()[0]:\n",
    "        ind_labels = np.hstack((ind_labels, val_point_labels[val_borders[ind]:val_borders[ind+1]]))\n",
    "    for ind in (test_labels == i).nonzero()[0]:\n",
    "        ind_labels = np.hstack((ind_labels, test_point_labels[test_borders[ind]:test_borders[ind+1]]))\n",
    "    label2point_labels[i] = np.unique(ind_labels)\n",
    "\n",
    "_, val_label2weight = np.unique(val_labels, return_counts=True)\n",
    "val_label2weight = np.float32(val_label2weight) / val_label2weight.sum()\n",
    "\n",
    "_, test_label2weight = np.unique(test_labels, return_counts=True)\n",
    "test_label2weight = np.float32(test_label2weight) / test_label2weight.sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import sys \n",
    "sys.setrecursionlimit(100000)\n",
    "\n",
    "import theano.sandbox.cuda\n",
    "theano.sandbox.cuda.use('gpu0')\n",
    "import theano\n",
    "import theano.tensor as T\n",
    "\n",
    "from lasagne.layers import InputLayer, ReshapeLayer, NonlinearityLayer, ExpressionLayer\n",
    "from lasagne.layers import ElemwiseSumLayer, ElemwiseMergeLayer, ConcatLayer\n",
    "from lasagne.layers import DenseLayer\n",
    "from lasagne.layers.dnn import BatchNormDNNLayer\n",
    "from lasagne.nonlinearities import rectify, softmax\n",
    "\n",
    "from lasagne.layers import get_output, get_all_params\n",
    "from lasagne.regularization import regularize_network_params, l2\n",
    "from lasagne.objectives import categorical_crossentropy, categorical_accuracy\n",
    "from lasagne.updates import adam\n",
    "\n",
    "from lib.nn.layers import SharedDotLayer, SPTNormReshapeLayer, SPTUpscaleLayer\n",
    "from lib.nn.utils import dump_weights, load_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "clouds = T.tensor3(dtype='float32')\n",
    "norms = [T.tensor3(dtype='float32') for step in xrange(config['steps'])]\n",
    "prob_masks = T.matrix(dtype='float32')\n",
    "target = T.matrix(dtype='int32')\n",
    "\n",
    "KDNet = {}\n",
    "if config['input_features'] == 'no':\n",
    "    KDNet['input'] = InputLayer((None, 1, 2**config['steps']), input_var=clouds)\n",
    "else:\n",
    "    KDNet['input'] = InputLayer((None, 3, 2**config['steps']), input_var=clouds)\n",
    "for i in xrange(config['steps']):\n",
    "    KDNet['norm{}_r'.format(i+1)] = InputLayer((None, 3, 2**(config['steps']-1-i)), input_var=norms[i])\n",
    "    KDNet['norm{}_l'.format(i+1)] = ExpressionLayer(KDNet['norm{}_r'.format(i+1)], lambda X: -X)\n",
    "    \n",
    "    KDNet['norm{}_r_X-'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(i+1)], '-', 0, config['n_f'][i+1])\n",
    "    KDNet['norm{}_r_Y-'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(i+1)], '-', 1, config['n_f'][i+1])\n",
    "    KDNet['norm{}_r_Z-'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(i+1)], '-', 2, config['n_f'][i+1])\n",
    "    KDNet['norm{}_r_X+'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(i+1)], '+', 0, config['n_f'][i+1])\n",
    "    KDNet['norm{}_r_Y+'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(i+1)], '+', 1, config['n_f'][i+1])\n",
    "    KDNet['norm{}_r_Z+'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(i+1)], '+', 2, config['n_f'][i+1])   \n",
    "    KDNet['norm{}_l_X-'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(i+1)], '-', 0, config['n_f'][i+1])\n",
    "    KDNet['norm{}_l_Y-'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(i+1)], '-', 1, config['n_f'][i+1])\n",
    "    KDNet['norm{}_l_Z-'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(i+1)], '-', 2, config['n_f'][i+1])\n",
    "    KDNet['norm{}_l_X+'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(i+1)], '+', 0, config['n_f'][i+1])\n",
    "    KDNet['norm{}_l_Y+'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(i+1)], '+', 1, config['n_f'][i+1])\n",
    "    KDNet['norm{}_l_Z+'.format(i+1)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(i+1)], '+', 2, config['n_f'][i+1])\n",
    "    \n",
    "    KDNet['cloud{}'.format(i+1)] = SharedDotLayer(KDNet['input'], config['n_f'][i]) if i == 0 else \\\n",
    "                                                  ElemwiseSumLayer([KDNet['cloud{}_l_X-_masked'.format(i)],\n",
    "                                                                    KDNet['cloud{}_l_Y-_masked'.format(i)],\n",
    "                                                                    KDNet['cloud{}_l_Z-_masked'.format(i)],\n",
    "                                                                    KDNet['cloud{}_l_X+_masked'.format(i)],\n",
    "                                                                    KDNet['cloud{}_l_Y+_masked'.format(i)],\n",
    "                                                                    KDNet['cloud{}_l_Z+_masked'.format(i)],\n",
    "                                                                    KDNet['cloud{}_r_X-_masked'.format(i)],\n",
    "                                                                    KDNet['cloud{}_r_Y-_masked'.format(i)],\n",
    "                                                                    KDNet['cloud{}_r_Z-_masked'.format(i)],\n",
    "                                                                    KDNet['cloud{}_r_X+_masked'.format(i)],\n",
    "                                                                    KDNet['cloud{}_r_Y+_masked'.format(i)],\n",
    "                                                                    KDNet['cloud{}_r_Z+_masked'.format(i)]])\n",
    "    KDNet['cloud{}_bn'.format(i+1)] = BatchNormDNNLayer(KDNet['cloud{}'.format(i+1)])\n",
    "    KDNet['cloud{}_relu'.format(i+1)] = NonlinearityLayer(KDNet['cloud{}_bn'.format(i+1)], rectify)\n",
    "    \n",
    "    KDNet['cloud{}_r'.format(i+1)] = ExpressionLayer(KDNet['cloud{}_relu'.format(i+1)],\n",
    "                                                     lambda X: X[:, :, 1::2], \n",
    "                                                     (None, config['n_f'][i], 2**(config['steps']-i-1)))\n",
    "    KDNet['cloud{}_l'.format(i+1)] = ExpressionLayer(KDNet['cloud{}_relu'.format(i+1)],\n",
    "                                                     lambda X: X[:, :, ::2], \n",
    "                                                     (None, config['n_f'][i], 2**(config['steps']-i-1)))\n",
    "    \n",
    "    KDNet['cloud{}_r_X-'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_r'.format(i+1)], config['n_f'][i+1])\n",
    "    KDNet['cloud{}_r_Y-'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_r'.format(i+1)], config['n_f'][i+1])\n",
    "    KDNet['cloud{}_r_Z-'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_r'.format(i+1)], config['n_f'][i+1])\n",
    "    KDNet['cloud{}_r_X+'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_r'.format(i+1)], config['n_f'][i+1])\n",
    "    KDNet['cloud{}_r_Y+'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_r'.format(i+1)], config['n_f'][i+1])\n",
    "    KDNet['cloud{}_r_Z+'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_r'.format(i+1)], config['n_f'][i+1])\n",
    "    KDNet['cloud{}_l_X-'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_l'.format(i+1)], config['n_f'][i+1],\n",
    "                                                       W=KDNet['cloud{}_r_X-'.format(i+1)].W,\n",
    "                                                       b=KDNet['cloud{}_r_X-'.format(i+1)].b)\n",
    "    KDNet['cloud{}_l_Y-'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_l'.format(i+1)], config['n_f'][i+1],\n",
    "                                                       W=KDNet['cloud{}_r_Y-'.format(i+1)].W,\n",
    "                                                       b=KDNet['cloud{}_r_Y-'.format(i+1)].b)\n",
    "    KDNet['cloud{}_l_Z-'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_l'.format(i+1)], config['n_f'][i+1],\n",
    "                                                       W=KDNet['cloud{}_r_Z-'.format(i+1)].W,\n",
    "                                                       b=KDNet['cloud{}_r_Z-'.format(i+1)].b)\n",
    "    KDNet['cloud{}_l_X+'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_l'.format(i+1)], config['n_f'][i+1],\n",
    "                                                       W=KDNet['cloud{}_r_X+'.format(i+1)].W,\n",
    "                                                       b=KDNet['cloud{}_r_X+'.format(i+1)].b)\n",
    "    KDNet['cloud{}_l_Y+'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_l'.format(i+1)], config['n_f'][i+1],\n",
    "                                                       W=KDNet['cloud{}_r_Y+'.format(i+1)].W,\n",
    "                                                       b=KDNet['cloud{}_r_Y+'.format(i+1)].b)\n",
    "    KDNet['cloud{}_l_Z+'.format(i+1)] = SharedDotLayer(KDNet['cloud{}_l'.format(i+1)], config['n_f'][i+1],\n",
    "                                                       W=KDNet['cloud{}_r_Z+'.format(i+1)].W,\n",
    "                                                       b=KDNet['cloud{}_r_Z+'.format(i+1)].b)\n",
    "    \n",
    "    KDNet['cloud{}_r_X-_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_r_X-'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_r_X-'.format(i+1)]], T.mul)\n",
    "    KDNet['cloud{}_r_Y-_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_r_Y-'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_r_Y-'.format(i+1)]], T.mul)\n",
    "    KDNet['cloud{}_r_Z-_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_r_Z-'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_r_Z-'.format(i+1)]], T.mul)\n",
    "    KDNet['cloud{}_r_X+_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_r_X+'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_r_X+'.format(i+1)]], T.mul)\n",
    "    KDNet['cloud{}_r_Y+_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_r_Y+'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_r_Y+'.format(i+1)]], T.mul)\n",
    "    KDNet['cloud{}_r_Z+_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_r_Z+'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_r_Z+'.format(i+1)]], T.mul)\n",
    "    KDNet['cloud{}_l_X-_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_l_X-'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_l_X-'.format(i+1)]], T.mul)\n",
    "    KDNet['cloud{}_l_Y-_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_l_Y-'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_l_Y-'.format(i+1)]], T.mul)\n",
    "    KDNet['cloud{}_l_Z-_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_l_Z-'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_l_Z-'.format(i+1)]], T.mul)\n",
    "    KDNet['cloud{}_l_X+_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_l_X+'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_l_X+'.format(i+1)]], T.mul)\n",
    "    KDNet['cloud{}_l_Y+_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_l_Y+'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_l_Y+'.format(i+1)]], T.mul)\n",
    "    KDNet['cloud{}_l_Z+_masked'.format(i+1)] = ElemwiseMergeLayer([KDNet['cloud{}_l_Z+'.format(i+1)],\n",
    "                                                                   KDNet['norm{}_l_Z+'.format(i+1)]], T.mul)\n",
    "    \n",
    "KDNet['cloud_fin'] = ElemwiseSumLayer([KDNet['cloud{}_l_X-_masked'.format(config['steps'])],\n",
    "                                       KDNet['cloud{}_l_Y-_masked'.format(config['steps'])],\n",
    "                                       KDNet['cloud{}_l_Z-_masked'.format(config['steps'])],\n",
    "                                       KDNet['cloud{}_l_X+_masked'.format(config['steps'])],\n",
    "                                       KDNet['cloud{}_l_Y+_masked'.format(config['steps'])],\n",
    "                                       KDNet['cloud{}_l_Z+_masked'.format(config['steps'])],\n",
    "                                       KDNet['cloud{}_r_X-_masked'.format(config['steps'])],\n",
    "                                       KDNet['cloud{}_r_Y-_masked'.format(config['steps'])],\n",
    "                                       KDNet['cloud{}_r_Z-_masked'.format(config['steps'])],\n",
    "                                       KDNet['cloud{}_r_X+_masked'.format(config['steps'])],\n",
    "                                       KDNet['cloud{}_r_Y+_masked'.format(config['steps'])],\n",
    "                                       KDNet['cloud{}_r_Z+_masked'.format(config['steps'])]])\n",
    "KDNet['cloud_fin_bn'] = BatchNormDNNLayer(KDNet['cloud_fin'])\n",
    "KDNet['cloud_fin_relu'] = NonlinearityLayer(KDNet['cloud_fin_bn'], rectify)\n",
    "KDNet['cloud_fin_reshape'] = ReshapeLayer(KDNet['cloud_fin_relu'], (-1, config['n_f'][-2]))\n",
    "\n",
    "KDNet['fin'] = DenseLayer(KDNet['cloud_fin_reshape'], config['n_f'][-1])\n",
    "\n",
    "KDNet['inv_cloud_fin_raw_reshape'] = DenseLayer(KDNet['fin'], config['n_f'][-2], nonlinearity=None)\n",
    "KDNet['inv_cloud_fin_raw'] = ReshapeLayer(KDNet['inv_cloud_fin_raw_reshape'], (-1, config['n_f'][-2], 1))\n",
    "KDNet['inv_cloud_fin_raw_bn'] = BatchNormDNNLayer(KDNet['inv_cloud_fin_raw'])\n",
    "KDNet['inv_cloud_fin_raw_relu'] = NonlinearityLayer(KDNet['inv_cloud_fin_raw_bn'], rectify)\n",
    "\n",
    "KDNet['add_cloud_fin'] = SharedDotLayer(KDNet['cloud_fin_relu'], config['n_f'][-2])\n",
    "KDNet['add_cloud_fin_bn'] = BatchNormDNNLayer(KDNet['add_cloud_fin'])\n",
    "KDNet['add_cloud_fin_relu'] = NonlinearityLayer(KDNet['add_cloud_fin_bn'], rectify)\n",
    "\n",
    "KDNet['inv_cloud_fin'] = ConcatLayer([KDNet['inv_cloud_fin_raw_relu'], \n",
    "                                      KDNet['add_cloud_fin_relu']], axis=1)\n",
    "\n",
    "for i in xrange(config['steps']):\n",
    "    KDNet['inv_norm{}_r_X-'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(config['steps']-i)], \n",
    "                                                                             '-', 0, config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_norm{}_r_Y-'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(config['steps']-i)], \n",
    "                                                                             '-', 1, config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_norm{}_r_Z-'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(config['steps']-i)], \n",
    "                                                                             '-', 2, config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_norm{}_r_X+'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(config['steps']-i)], \n",
    "                                                                             '+', 0, config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_norm{}_r_Y+'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(config['steps']-i)], \n",
    "                                                                             '+', 1, config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_norm{}_r_Z+'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_r'.format(config['steps']-i)], \n",
    "                                                                             '+', 2, config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_norm{}_l_X-'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(config['steps']-i)], \n",
    "                                                                             '-', 0, config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_norm{}_l_Y-'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(config['steps']-i)], \n",
    "                                                                             '-', 1, config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_norm{}_l_Z-'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(config['steps']-i)], \n",
    "                                                                             '-', 2, config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_norm{}_l_X+'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(config['steps']-i)], \n",
    "                                                                             '+', 0, config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_norm{}_l_Y+'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(config['steps']-i)], \n",
    "                                                                             '+', 1, config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_norm{}_l_Z+'.format(config['steps']-i)] = SPTNormReshapeLayer(KDNet['norm{}_l'.format(config['steps']-i)], \n",
    "                                                                             '+', 2, config['n_f'][config['steps']-i-1])\n",
    "    \n",
    "    KDNet['inv_cloud{}_r_X-'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_cloud{}_r_Y-'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_cloud{}_r_Z-'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_cloud{}_r_X+'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_cloud{}_r_Y+'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_cloud{}_r_Z+'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1])\n",
    "    KDNet['inv_cloud{}_l_X-'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1],\n",
    "        W=KDNet['inv_cloud{}_r_X-'.format(config['steps']-i)].W, b=KDNet['inv_cloud{}_r_X-'.format(config['steps']-i)].b)\n",
    "    KDNet['inv_cloud{}_l_Y-'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1],\n",
    "        W=KDNet['inv_cloud{}_r_Y-'.format(config['steps']-i)].W, b=KDNet['inv_cloud{}_r_Y-'.format(config['steps']-i)].b)\n",
    "    KDNet['inv_cloud{}_l_Z-'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1],\n",
    "        W=KDNet['inv_cloud{}_r_Z-'.format(config['steps']-i)].W, b=KDNet['inv_cloud{}_r_Z-'.format(config['steps']-i)].b)\n",
    "    KDNet['inv_cloud{}_l_X+'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1],\n",
    "        W=KDNet['inv_cloud{}_r_X+'.format(config['steps']-i)].W, b=KDNet['inv_cloud{}_r_X+'.format(config['steps']-i)].b)\n",
    "    KDNet['inv_cloud{}_l_Y+'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1],\n",
    "        W=KDNet['inv_cloud{}_r_Y+'.format(config['steps']-i)].W, b=KDNet['inv_cloud{}_r_Y+'.format(config['steps']-i)].b)\n",
    "    KDNet['inv_cloud{}_l_Z+'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['inv_cloud{}'.format('_fin' if i == 0 else config['steps']-i+1)], config['n_f'][config['steps']-i-1],\n",
    "        W=KDNet['inv_cloud{}_r_Z+'.format(config['steps']-i)].W, b=KDNet['inv_cloud{}_r_Z+'.format(config['steps']-i)].b)\n",
    "    \n",
    "    KDNet['inv_cloud{}_r_X-_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_r_X-'.format(config['steps']-i)], \n",
    "                            KDNet['inv_norm{}_r_X-'.format(config['steps']-i)]], T.mul)\n",
    "    KDNet['inv_cloud{}_r_Y-_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_r_Y-'.format(config['steps']-i)],\n",
    "                            KDNet['inv_norm{}_r_Y-'.format(config['steps']-i)]], T.mul)\n",
    "    KDNet['inv_cloud{}_r_Z-_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_r_Z-'.format(config['steps']-i)],\n",
    "                            KDNet['inv_norm{}_r_Z-'.format(config['steps']-i)]], T.mul)\n",
    "    KDNet['inv_cloud{}_r_X+_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_r_X+'.format(config['steps']-i)], \n",
    "                            KDNet['inv_norm{}_r_X+'.format(config['steps']-i)]], T.mul)\n",
    "    KDNet['inv_cloud{}_r_Y+_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_r_Y+'.format(config['steps']-i)],\n",
    "                            KDNet['inv_norm{}_r_Y+'.format(config['steps']-i)]], T.mul)\n",
    "    KDNet['inv_cloud{}_r_Z+_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_r_Z+'.format(config['steps']-i)],\n",
    "                            KDNet['inv_norm{}_r_Z+'.format(config['steps']-i)]], T.mul)\n",
    "    KDNet['inv_cloud{}_l_X-_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_l_X-'.format(config['steps']-i)], \n",
    "                            KDNet['inv_norm{}_l_X-'.format(config['steps']-i)]], T.mul)\n",
    "    KDNet['inv_cloud{}_l_Y-_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_l_Y-'.format(config['steps']-i)],\n",
    "                            KDNet['inv_norm{}_l_Y-'.format(config['steps']-i)]], T.mul)\n",
    "    KDNet['inv_cloud{}_l_Z-_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_l_Z-'.format(config['steps']-i)],\n",
    "                            KDNet['inv_norm{}_l_Z-'.format(config['steps']-i)]], T.mul)\n",
    "    KDNet['inv_cloud{}_l_X+_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_l_X+'.format(config['steps']-i)], \n",
    "                            KDNet['inv_norm{}_l_X+'.format(config['steps']-i)]], T.mul)\n",
    "    KDNet['inv_cloud{}_l_Y+_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_l_Y+'.format(config['steps']-i)],\n",
    "                            KDNet['inv_norm{}_l_Y+'.format(config['steps']-i)]], T.mul)\n",
    "    KDNet['inv_cloud{}_l_Z+_masked'.format(config['steps']-i)] = \\\n",
    "        ElemwiseMergeLayer([KDNet['inv_cloud{}_l_Z+'.format(config['steps']-i)],\n",
    "                            KDNet['inv_norm{}_l_Z+'.format(config['steps']-i)]], T.mul)\n",
    "\n",
    "    KDNet['inv_cloud{}_r'.format(config['steps']-i)] = ElemwiseSumLayer(\n",
    "        [KDNet['inv_cloud{}_r_X-_masked'.format(config['steps']-i)],\n",
    "         KDNet['inv_cloud{}_r_Y-_masked'.format(config['steps']-i)],\n",
    "         KDNet['inv_cloud{}_r_Z-_masked'.format(config['steps']-i)],\n",
    "         KDNet['inv_cloud{}_r_X+_masked'.format(config['steps']-i)],\n",
    "         KDNet['inv_cloud{}_r_Y+_masked'.format(config['steps']-i)],\n",
    "         KDNet['inv_cloud{}_r_Z+_masked'.format(config['steps']-i)]])\n",
    "    KDNet['inv_cloud{}_l'.format(config['steps']-i)] = ElemwiseSumLayer(\n",
    "        [KDNet['inv_cloud{}_l_X-_masked'.format(config['steps']-i)],\n",
    "         KDNet['inv_cloud{}_l_Y-_masked'.format(config['steps']-i)],\n",
    "         KDNet['inv_cloud{}_l_Z-_masked'.format(config['steps']-i)],\n",
    "         KDNet['inv_cloud{}_l_X+_masked'.format(config['steps']-i)],\n",
    "         KDNet['inv_cloud{}_l_Y+_masked'.format(config['steps']-i)],\n",
    "         KDNet['inv_cloud{}_l_Z+_masked'.format(config['steps']-i)]])\n",
    "    \n",
    "    KDNet['inv_cloud{}_r_up'.format(config['steps']-i)] = SPTUpscaleLayer(\n",
    "        KDNet['inv_cloud{}_r'.format(config['steps']-i)], mode='right')\n",
    "    KDNet['inv_cloud{}_l_up'.format(config['steps']-i)] = SPTUpscaleLayer(\n",
    "        KDNet['inv_cloud{}_l'.format(config['steps']-i)], mode='left')\n",
    "    KDNet['inv_cloud{}_raw'.format(config['steps']-i)] = ElemwiseSumLayer(\n",
    "        [KDNet['inv_cloud{}_r_up'.format(config['steps']-i)], \n",
    "         KDNet['inv_cloud{}_l_up'.format(config['steps']-i)]])\n",
    "    KDNet['inv_cloud{}_raw_bn'.format(config['steps']-i)] = BatchNormDNNLayer(\n",
    "        KDNet['inv_cloud{}_raw'.format(config['steps']-i)])\n",
    "    KDNet['inv_cloud{}_raw_relu'.format(config['steps']-i)] = NonlinearityLayer(\n",
    "        KDNet['inv_cloud{}_raw_bn'.format(config['steps']-i)], rectify)\n",
    "    \n",
    "    KDNet['add_cloud{}'.format(config['steps']-i)] = SharedDotLayer(\n",
    "        KDNet['cloud{}_relu'.format(config['steps']-i)], config['n_f'][config['steps']-i-1])\n",
    "    KDNet['add_cloud{}_bn'.format(config['steps']-i)] = BatchNormDNNLayer(\n",
    "        KDNet['add_cloud{}'.format(config['steps']-i)])\n",
    "    KDNet['add_cloud{}_relu'.format(config['steps']-i)] = NonlinearityLayer(\n",
    "        KDNet['add_cloud{}_bn'.format(config['steps']-i)], rectify)\n",
    "    \n",
    "    KDNet['inv_cloud{}'.format(config['steps']-i)] = ConcatLayer(\n",
    "        [KDNet['inv_cloud{}_raw_relu'.format(config['steps']-i)], \n",
    "         KDNet['add_cloud{}_relu'.format(config['steps']-i)]], axis=1)\n",
    "\n",
    "for i in xrange(len(config['n_mlp'])):\n",
    "    KDNet['mlp{}'.format(i+1)] = SharedDotLayer(KDNet['inv_cloud1' if i == 0 else ('mlp'+str(i))], config['n_mlp'][i])\n",
    "    KDNet['mlp{}_bn'.format(i+1)] = BatchNormDNNLayer(KDNet['mlp{}'.format(i+1)])\n",
    "    KDNet['mlp{}_relu'.format(i+1)] = NonlinearityLayer(KDNet['mlp{}_bn'.format(i+1)], rectify)\n",
    "\n",
    "KDNet['final'] = SharedDotLayer(KDNet['mlp{}_relu'.format(len(config['n_mlp']))], config['n_output'])\n",
    "\n",
    "# Loading weights (optional)\n",
    "# load_weights('./models/RT+AS+TR+1e-4_10.pkl', KDNet['final'])\n",
    "\n",
    "final = get_output(KDNet['final'])\n",
    "final_exp = T.exp(final - (final*prob_masks.dimshuffle(0, 1, 'x')).max(axis=1, keepdims=True))\n",
    "final_exp_masked = final_exp*prob_masks.dimshuffle(0, 1, 'x')\n",
    "probs = final_exp_masked / final_exp_masked.sum(axis=1, keepdims=True)\n",
    "\n",
    "final_det = get_output(KDNet['final'], deterministic=True)\n",
    "final_det_exp = T.exp(final_det - (final_det*prob_masks.dimshuffle(0, 1, 'x')).max(axis=1, keepdims=True))\n",
    "final_det_exp_masked = final_det_exp*prob_masks.dimshuffle(0, 1, 'x')\n",
    "probs_det = final_det_exp_masked / final_det_exp_masked.sum(axis=1, keepdims=True)\n",
    "\n",
    "weights = get_all_params(KDNet['final'], trainable=True)\n",
    "l2_pen = regularize_network_params(KDNet['final'], l2)\n",
    "\n",
    "target_shaped = T.cast(T.eq(T.tile(np.arange(config['n_output'], dtype=np.int32).reshape(1, -1, 1), \n",
    "                                   (target.shape[0], 1, target.shape[1])),\n",
    "                            target.dimshuffle(0, 'x', 1)), \n",
    "                       'float32')\n",
    "\n",
    "loss = T.mean(-T.sum(target_shaped*T.log(probs + 1e-8), axis=1), axis=1).mean() + config['l2']*l2_pen\n",
    "acc = T.cast(T.eq(T.argmax(probs, axis=1), target), 'int32').sum()\n",
    "\n",
    "lr = theano.shared(np.float32(config['lr']))\n",
    "updates = adam(loss, weights, learning_rate=lr)\n",
    "\n",
    "train_fun = theano.function([clouds] + norms + [prob_masks, target], [loss, acc], updates=updates)\n",
    "prob_fun = theano.function([clouds] + norms + [prob_masks], probs_det)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build Data Iterator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from copy import copy\n",
    "\n",
    "from lib.generators.pointcloud import generate_clouds\n",
    "from lib.trees.kdtrees import KDTrees"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def Leaves2Points(cloud, idx, normals, medians, steps=10):\n",
    "    mask0 = [np.arange(cloud.shape[1])]\n",
    "    \n",
    "    for step in xrange(steps):\n",
    "        mask1 = []\n",
    "        for i, submask in enumerate(mask0):\n",
    "            tmp_cloud = cloud[:, submask]\n",
    "            mask_buf = medians[step][idx, i] <= (normals[step][idx, :, i].reshape(-1, 1)*tmp_cloud).sum(axis=0)\n",
    "            mask1.append(submask[np.logical_not(mask_buf).nonzero()[0]])\n",
    "            mask1.append(submask[mask_buf.nonzero()[0]])\n",
    "        mask0 = copy(mask1)\n",
    "        \n",
    "    return mask0\n",
    "\n",
    "\n",
    "def iterate_minibatches(*arrays, **kwargs):\n",
    "    if kwargs['mode'] == 'train':\n",
    "        indices = np.random.choice((len(arrays[1]) - 1), \n",
    "                                   size=(len(arrays[1]) - 1)/kwargs['batchsize']*kwargs['batchsize'])\n",
    "    elif kwargs['mode'] == 'test':\n",
    "        indices = np.arange(len(arrays[1]) - 1)\n",
    "    if kwargs['shuffle']:\n",
    "        np.random.shuffle(indices)\n",
    "        \n",
    "    for start_idx in xrange(0, len(indices), kwargs['batchsize']):\n",
    "        excerpt = indices[start_idx:start_idx + kwargs['batchsize']]\n",
    "        tmp, smpl = generate_clouds(excerpt, kwargs['steps'], arrays[0], arrays[1])\n",
    "        labels = arrays[2][excerpt]\n",
    "        \n",
    "        prob_masks = np.zeros((len(excerpt), kwargs['n_output']), dtype=np.float32)\n",
    "        for i, label in enumerate(labels):\n",
    "            prob_masks[i][label2point_labels[label]] = 1.\n",
    "        \n",
    "        if kwargs['flip']:\n",
    "            flip = np.random.random(size=(len(tmp), 2, 1))\n",
    "            flip[flip >= 0.5] = 1.\n",
    "            flip[flip < 0.5] = -1.\n",
    "            tmp[:, :2] *= flip\n",
    "        \n",
    "        if kwargs['ascale']:\n",
    "            ascale = (kwargs['as_min'] + (kwargs['as_max'] - kwargs['as_min'])*np.random.random(size=(len(tmp), kwargs['dim'], 1)))\n",
    "            tmp *= ascale\n",
    "            tmp /= np.fabs(tmp).max(axis=(1, 2), keepdims=True)\n",
    "            \n",
    "        if kwargs['rotate']:\n",
    "            r = np.sqrt((tmp[:, :2]**2).sum(axis=1))\n",
    "            coss = tmp[:, 0]/r\n",
    "            sins = tmp[:, 1]/r\n",
    "            \n",
    "            if kwargs['test_pos'] is not None:\n",
    "                alpha = 2*np.pi*kwargs['test_pos']/kwargs['r_positions']\n",
    "            else:\n",
    "                alpha = 2*np.pi*np.random.randint(0, kwargs['r_positions'], (len(tmp), 1))/kwargs['positions']\n",
    "                \n",
    "            cosr = np.cos(alpha)\n",
    "            sinr = np.sin(alpha)\n",
    "            cos = coss*cosr - sins*sinr\n",
    "            sin = sins*cosr + sinr*coss\n",
    "            tmp[:, 0] = r*cos\n",
    "            tmp[:, 1] = r*sin\n",
    "            \n",
    "        if kwargs['translate']:\n",
    "            mins = tmp.min(axis=2, keepdims=True)\n",
    "            maxs = tmp.max(axis=2, keepdims=True)\n",
    "            rngs = maxs - mins\n",
    "            translate = kwargs['t_rate']*(np.random.random(size=(len(tmp), kwargs['dim'], 1)) - 0.5)*rngs\n",
    "            tmp += translate\n",
    "\n",
    "        trees_data = KDTrees(tmp, dim=kwargs['dim'], steps=kwargs['steps'], \n",
    "                             lim=kwargs['lim'], det=kwargs['det'], gamma=kwargs['gamma'],\n",
    "                             medians=False if kwargs['mode'] == 'train' else True)\n",
    "            \n",
    "        sortings, normals, medians = trees_data['sortings'], trees_data['normals'], trees_data.get('medians', None)\n",
    "        \n",
    "        clouds = np.empty((len(excerpt), 3 if kwargs['input_features'] == 'all' else 1, 2**kwargs['steps']), \n",
    "                          dtype=np.float32)\n",
    "        \n",
    "        if kwargs['mode'] == 'train':\n",
    "            tmp_point_labels = arrays[3][smpl]\n",
    "            clouds_point_labels = np.empty((len(excerpt), 2**kwargs['steps']), dtype=np.int32)\n",
    "            for i, srt in enumerate(sortings):\n",
    "                clouds[i] = tmp[i, :, srt].T\n",
    "                clouds_point_labels[i] = tmp_point_labels[i, srt]\n",
    "                \n",
    "            yield [clouds] + normals[::-1] + [prob_masks, clouds_point_labels]\n",
    "            \n",
    "        if kwargs['mode'] == 'test':\n",
    "            for i, srt in enumerate(sortings):\n",
    "                clouds[i] = tmp[i, :, srt].T\n",
    "            \n",
    "            mappings = []\n",
    "            for i, idx in enumerate(excerpt):\n",
    "                pcoord4mapping = np.transpose(arrays[0][arrays[1][idx]:arrays[1][idx+1], :].copy(), (1, 0))\n",
    "                \n",
    "                if kwargs['flip']:\n",
    "                    pcoord4mapping[:2] *= flip[i]\n",
    "                \n",
    "                if kwargs['ascale']:\n",
    "                    pcoord4mapping *= ascale[i]\n",
    "                    pcoord4mapping /= np.fabs(pcoord4mapping).max()\n",
    "                    \n",
    "                if kwargs['translate']:\n",
    "                    pcoord4mapping += translate[i]\n",
    "                        \n",
    "                mappings.append(\n",
    "                    Leaves2Points(pcoord4mapping.astype(np.float32), i, normals, medians, steps=kwargs['steps']))\n",
    "            \n",
    "            yield [clouds] + normals[::-1] + [prob_masks, mappings, excerpt]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prediction function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from sklearn.metrics import confusion_matrix\n",
    "\n",
    "\n",
    "def get_probs(coordinates, borders, labels, **kwargs):\n",
    "    predictions = []\n",
    "    for i in xrange(len(borders)-1):\n",
    "        predictions.append(np.zeros((borders[i+1]-borders[i], kwargs['n_output']), dtype=np.float32))\n",
    "        \n",
    "    for ens in xrange(kwargs['n_ens']):\n",
    "        for i, batch in enumerate(iterate_minibatches(coordinates, borders, labels, **kwargs)):\n",
    "            probs = prob_fun(batch[0], \n",
    "                             batch[1], batch[2], batch[3], batch[4], batch[5], batch[6], \n",
    "                             batch[7], batch[8], batch[9], batch[10], batch[11], batch[12],\n",
    "                             batch[13])\n",
    "            \n",
    "            for j, mapping in enumerate(batch[-2]):\n",
    "                for k, inds in enumerate(mapping):\n",
    "                    predictions[batch[-1][j]][inds] += probs[j, :, k].reshape(1, -1)\n",
    "        \n",
    "    return list(map(lambda pred: pred/kwargs['n_ens'], predictions))\n",
    "\n",
    "\n",
    "def acc_fun(borders, point_labels, probs):\n",
    "    n_true = 0.\n",
    "    n_all = 0.\n",
    "    \n",
    "    for i, prob in enumerate(probs):\n",
    "        tmp_point_labels = point_labels[borders[i]:borders[i+1]]\n",
    "        n_true += (prob.argmax(1) == tmp_point_labels).sum()\n",
    "        n_all += len(tmp_point_labels)\n",
    "        \n",
    "    return 100.*n_true/n_all\n",
    "\n",
    "\n",
    "def label_acc_fun(borders, labels, point_labels, probs):\n",
    "    label2n_true = np.zeros(16, dtype=np.int32)\n",
    "    label2n_all = np.zeros(16, dtype=np.int32)\n",
    "        \n",
    "    for i, prob in enumerate(probs):\n",
    "        label2n_true[labels[i]] += (prob.argmax(1) == point_labels[borders[i]:borders[i+1]]).sum()\n",
    "        label2n_all[labels[i]] += len(prob)\n",
    "        \n",
    "    return 100. * label2n_true / label2n_all\n",
    "\n",
    "\n",
    "def mIoU_fun(borders, labels, point_labels, probs):\n",
    "    sIoU = np.zeros(16, dtype=np.float32)\n",
    "    _, label_counters = np.unique(labels, return_counts=True)\n",
    "    \n",
    "    for i, prob in enumerate(probs):\n",
    "        C = confusion_matrix(point_labels[borders[i]:borders[i+1]], prob.argmax(1), \n",
    "                             labels=label2point_labels[labels[i]])\n",
    "        \n",
    "        intersections = C[np.arange(C.shape[0]), np.arange(C.shape[0])]\n",
    "        unions = (C.sum(axis=0) + C.sum(axis=1) - intersections)\n",
    "        \n",
    "        sIoU_buf = 0.\n",
    "        for j, union in enumerate(unions):\n",
    "            if union == 0:\n",
    "                sIoU_buf += 100.\n",
    "            else:\n",
    "                sIoU_buf += 100.*np.float32(intersections[j])/union\n",
    "                \n",
    "        sIoU[labels[i]] += sIoU_buf/C.shape[0]\n",
    "        \n",
    "    return sIoU / label_counters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from sys import stdout\n",
    "from time import time\n",
    "\n",
    "config['mode'] = 'train'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "num_epochs = 100\n",
    "num_save = 5\n",
    "\n",
    "# lr.set_value(np.float32(1e-4))\n",
    "\n",
    "for epoch in xrange(num_epochs):\n",
    "    train_err = 0.\n",
    "    train_acc = 0.\n",
    "    \n",
    "    start_time = time()\n",
    "    for i, batch in enumerate(iterate_minibatches(train_points, train_borders, \n",
    "                                                  train_labels, train_point_labels, **config)):\n",
    "        train_err_batch, train_acc_batch = train_fun(batch[0], \n",
    "                                                     batch[1], batch[2], batch[3], batch[4], batch[5], batch[6], \n",
    "                                                     batch[7], batch[8], batch[9], batch[10], batch[11], batch[12],\n",
    "                                                     batch[13],\n",
    "                                                     batch[14])\n",
    "        train_err += train_err_batch*len(batch[0])\n",
    "        train_acc += train_acc_batch\n",
    "\n",
    "        stdout.write('\\rEpoch progress: {}/{}\\tAccuracy: {:.2f} %\\t\\tLoss: {:.5f}'\n",
    "                     .format(config['batchsize']*(i+1),\n",
    "                     len(train_labels)/config['batchsize']*config['batchsize'],\n",
    "                     100*train_acc/(2**config['steps'])/(config['batchsize']*(i+1)),\n",
    "                     train_err/(config['batchsize']*(i+1))))\n",
    "        stdout.flush()\n",
    "    stdout.write('\\n')\n",
    "    stdout.flush()\n",
    "        \n",
    "    print(\"Epoch {} of {} took {:.3f}s\".format(epoch + 1, num_epochs, time() - start_time))\n",
    "    print(\"  training loss (in-iteration):\\t\\t{:.6f}\".format(train_err/(config['batchsize']*(i+1))))\n",
    "    print(\"  train accuracy:\\t\\t{:.2f} %\".format(100*train_acc/(2**config['steps'])/(config['batchsize']*(i+1))))\n",
    "\n",
    "    if (epoch+1) % (num_epochs/num_save) == 0:\n",
    "        config['mode'] = 'test'\n",
    "        val_probs = get_probs(val_points, val_borders, val_labels, **config)\n",
    "        val_acc = acc_fun(val_borders, val_point_labels, val_probs)\n",
    "        val_label2acc = label_acc_fun(val_borders, val_labels, val_point_labels, val_probs)\n",
    "        val_mIoU = mIoU_fun(val_borders, val_labels, val_point_labels, val_probs)\n",
    "        print(\"  val accuracy:\\t\\t\\t{:.2f}%\".format(val_acc))\n",
    "        print(\"  val category accuracy:\")\n",
    "        for j in xrange(16):\n",
    "            print(\"  \\t{}\\t\\t{:.2f}%\".format(label2name[j], val_label2acc[j]))\n",
    "        print(\"  \\tMEAN\\t\\t\\t{:.2f}%\".format(val_label2acc.mean()))\n",
    "        print(\"  val mIoU:\")\n",
    "        for j in xrange(16):\n",
    "            print(\"  \\t{}\\t\\t{:.2f}%\".format(label2name[j], val_mIoU[j]))\n",
    "        print(\"  \\tMEAN\\t\\t\\t{:.2f}%\".format((val_mIoU*val_label2weight).sum()))\n",
    "        config['mode'] = 'train'\n",
    "        \n",
    "        dump_weights('./models/RT_AS+TR_10.pkl', KDNet['final'])\n",
    "        print '  Model saved!'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "Python [default]",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}