docs/demos/xfer-custom-repurposers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Writing a custom Repurposer\n",
    "\n",
    "Xfer implements and supports two kinds of Repurposers:\n",
    "\n",
    "\n",
    "* **Meta-model Repurposer** - this uses the source model to extract features and then fits a meta-model to the features\n",
    "* **Neural network Repurposer** - this modifies the source model to create a target model\n",
    "\n",
    "Below are examples of creating custom Repurposers for both classes\n",
    "\n",
    "### Setup\n",
    "First import relevant modules, define data iterators and load a source model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "import logging\n",
    "logging.disable(logging.WARNING)\n",
    "\n",
    "import xfer\n",
    "\n",
    "import os\n",
    "import glob\n",
    "import mxnet as mx\n",
    "import random\n",
    "from sklearn.model_selection import StratifiedShuffleSplit\n",
    "from sklearn.metrics import classification_report\n",
    "\n",
    "random.seed(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def get_iterators_from_folder(data_dir, train_size=0.6, batchsize=10, label_name='softmax_label', data_name='data', random_state=1):\n",
    "    \"\"\"\n",
    "    Method to create iterators from data stored in a folder with the following structure:\n",
    "    /data_dir\n",
    "        /class1\n",
    "            class1_img1\n",
    "            class1_img2\n",
    "            ...\n",
    "            class1_imgN\n",
    "        /class2\n",
    "            class2_img1\n",
    "            class2_img2\n",
    "            ...\n",
    "            class2_imgN\n",
    "        ...\n",
    "        /classN\n",
    "    \"\"\"\n",
    "    # assert dir exists\n",
    "    if not os.path.isdir(data_dir):\n",
    "        raise ValueError('Directory not found: {}'.format(data_dir))\n",
    "    # get class names\n",
    "    classes = [x.split('/')[-1] for x in glob.glob(data_dir+'/*')]\n",
    "    classes.sort()\n",
    "    fnames = []\n",
    "    labels = []\n",
    "    for c in classes:\n",
    "            # get all the image filenames and labels\n",
    "            images = glob.glob(data_dir+'/'+c+'/*')\n",
    "            images.sort()\n",
    "            fnames += images\n",
    "            labels += [c]*len(images)\n",
    "    # create label2id mapping\n",
    "    id2label = dict(enumerate(set(labels)))\n",
    "    label2id = dict((v,k) for k, v in id2label.items())\n",
    "\n",
    "    # get indices of train and test\n",
    "    sss = StratifiedShuffleSplit(n_splits=2, test_size=None, train_size=train_size, random_state=random_state)\n",
    "    train_indices, test_indices = next(sss.split(labels, labels))\n",
    "    \n",
    "    train_img_list = []\n",
    "    test_img_list = []\n",
    "    train_labels = []\n",
    "    test_labels = []\n",
    "    # create imglist for training and test\n",
    "    for idx in train_indices:\n",
    "        train_img_list.append([label2id[labels[idx]], fnames[idx]])\n",
    "        train_labels.append(label2id[labels[idx]])\n",
    "    for idx in test_indices:\n",
    "        test_img_list.append([label2id[labels[idx]], fnames[idx]])\n",
    "        test_labels.append(label2id[labels[idx]])\n",
    "        \n",
    "    # make iterators\n",
    "    train_iterator = mx.image.ImageIter(batchsize, (3,224,224), imglist=train_img_list, label_name=label_name, data_name=data_name,\n",
    "                                        path_root='')\n",
    "    test_iterator = mx.image.ImageIter(batchsize, (3,224,224), imglist=test_img_list, label_name=label_name, data_name=data_name,\n",
    "                                      path_root='')\n",
    "\n",
    "    return train_iterator, test_iterator, train_labels, test_labels, id2label, label2id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "dataset = 'test_images' # options are: 'test_sketches', 'test_images_sketch', 'mnist-50', 'test_images' or your own data.\n",
    "num_classes = 4\n",
    "\n",
    "train_iterator, test_iterator, train_labels, test_labels, id2label, label2id = get_iterators_from_folder(dataset, 0.6, 4, label_name='prob_label', random_state=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['vgg19-0000.params', 'vgg19-symbol.json']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Download vgg19 (trained on imagenet)\n",
    "path = 'http://data.mxnet.io/models/imagenet/'\n",
    "[mx.test_utils.download(path+'vgg/vgg19-0000.params'),\n",
    "mx.test_utils.download(path+'vgg/vgg19-symbol.json')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# This will be the source model we use for repurposing later\n",
    "source_model = mx.module.Module.load('vgg19', 0, label_names=['prob_label'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Meta-model Repurposer\n",
    "\n",
    "We will create a new Repurposer that uses the KNN algorithm as a meta-model.  The resulting Meta-model Repurposer will classify the features extracted by the neural network source model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from sklearn.neighbors import KNeighborsClassifier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class KNNRepurposer(xfer.MetaModelRepurposer):\n",
    "    def __init__(self, source_model: mx.mod.Module, feature_layer_names, context_function=mx.context.cpu, num_devices=1,\n",
    "                 n_neighbors=5, weights='uniform', algorithm='auto', leaf_size=30, p=2, metric='minkowski', metric_params=None, n_jobs=-1):\n",
    "        # Call init() of parent\n",
    "        super(KNNRepurposer, self).__init__(source_model, feature_layer_names, context_function, num_devices)\n",
    "        \n",
    "        # Initialise parameters specific to the KNN algorithm\n",
    "        self.n_neighbors = n_neighbors\n",
    "        self.weights = weights\n",
    "        self.algorithm = algorithm\n",
    "        self.leaf_size = leaf_size\n",
    "        self.p = p\n",
    "        self.metric = metric\n",
    "        self.metric_params = metric_params\n",
    "        self.n_jobs = n_jobs\n",
    "        \n",
    "    # Define function that takes a set of features and labels and returns a trained model.\n",
    "    # feature_indices_per_layer is a dictionary which gives the feature indices which correspond\n",
    "    # to each layer's features.\n",
    "    def _train_model_from_features(self, features, labels, feature_indices_per_layer=None):\n",
    "        lin_model = KNeighborsClassifier(n_neighbors=self.n_neighbors,\n",
    "                                        weights=self.weights,\n",
    "                                        algorithm=self.algorithm, \n",
    "                                        leaf_size=self.leaf_size, \n",
    "                                        p=self.p,\n",
    "                                        metric=self.metric,\n",
    "                                        metric_params=self.metric_params)\n",
    "        lin_model.fit(features, labels)\n",
    "        return lin_model\n",
    "    \n",
    "    # Define a function that predicts the class probability given features\n",
    "    def _predict_probability_from_features(self, features):\n",
    "        return self.target_model.predict_proba(features)\n",
    "    \n",
    "    # Define a function that predicts the class label given features\n",
    "    def _predict_label_from_features(self, features):\n",
    "        return self.target_model.predict(features)\n",
    "    \n",
    "    # In order to make your repurposer serialisable, you will need to implement functions\n",
    "    # which convert your model's parameters to a dictionary.\n",
    "    def get_params(self):\n",
    "        \"\"\"\n",
    "        This function should return a dictionary of all the parameters of the repurposer that\n",
    "        are in the repurposer constructor arguments.\n",
    "        \"\"\"\n",
    "        param_dict = super().get_params()\n",
    "        param_dict['n_neighbors'] = self.n_neighbors\n",
    "        param_dict['weights'] = self.weights \n",
    "        param_dict['algorithm'] = self.algorithm\n",
    "        param_dict['leaf_size'] = self.leaf_size\n",
    "        param_dict['p'] = self.p\n",
    "        param_dict['metric'] = self.metric \n",
    "        param_dict['metric_params'] = self.metric_params\n",
    "        param_dict['n_jobs'] = self.n_jobs\n",
    "        return param_dict\n",
    "    \n",
    "    # Some repurposers will need a get_attributes() and set_attributes() to get and set the parameters\n",
    "    # of the repurposer that are not in the constructor argument. An example is shown below:\n",
    "    \n",
    "    # def get_attributes(self):\n",
    "    #     \"\"\"\n",
    "    #     This function should return a dictionary of all the parameters of the repurposer that\n",
    "    #     are NOT in the constructor arguments.\n",
    "    #     \n",
    "    #     This function does not need to be defined if the repurposer has no specific attributes.\n",
    "    #     \"\"\"\n",
    "    #     param_dict = super().get_attributes()\n",
    "    #     param_dict['example_attribute'] = self.example_attribute\n",
    "    #     return param_dict\n",
    "    \n",
    "    # def set_attributes(self, input_dict):\n",
    "    #     super().set_attributes(input_dict)\n",
    "    #     self.example_attribute  = input_dict['example_attribute']\n",
    "    \n",
    "    def serialize(self, file_prefix):\n",
    "        \"\"\"\n",
    "        Saves repurposer (excluding source model) to file_prefix.json.\n",
    "        This method converts the repurposer to dictionary and saves as a json.\n",
    "\n",
    "\n",
    "        :param str file_prefix: Prefix to save file with\n",
    "        \"\"\"\n",
    "        output_dict = {}\n",
    "        output_dict[repurposer_keys.PARAMS] = self.get_params()\n",
    "        output_dict[repurposer_keys.TARGET_MODEL] = target_model_to_dict()  # This should be some serialised representation of the target model\n",
    "        output_dict.update(self.get_attributes())\n",
    "\n",
    "        utils.save_json(file_prefix, output_dict)\n",
    "\n",
    "    def deserialize(self, input_dict):\n",
    "        \"\"\"\n",
    "        Uses dictionary to set attributes of repurposer\n",
    "\n",
    "        :param dict input_dict: Dictionary containing values for attributes to be set to\n",
    "        \"\"\"\n",
    "        self.set_attributes(input_dict)  # Set attributes of the repurposer from input_dict\n",
    "        self.target_model = target_model_from_dict()  # Unpack dictionary representation of target model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "repurposerKNN = KNNRepurposer(source_model, ['fc8'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "repurposerKNN.repurpose(train_iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "results = repurposerKNN.predict_label(test_iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "             precision    recall  f1-score   support\n",
      "\n",
      "          0       1.00      0.50      0.67         2\n",
      "          1       0.67      1.00      0.80         2\n",
      "          2       1.00      1.00      1.00         2\n",
      "          3       1.00      1.00      1.00         2\n",
      "\n",
      "avg / total       0.92      0.88      0.87         8\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_report(y_pred=results, y_true=test_labels))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Neural Network Repurposer\n",
    "\n",
    "Now we will define a custom Neural Network Repurposer which performs transfer learning by:\n",
    "\n",
    "1. taking the original source neural network and keeping all layers up to `transfer_layer_name`\n",
    "2. adding two fully connected layers on the top\n",
    "3. fine-tuning with any conv layers frozen\n",
    "\n",
    "#### Definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Add2FullyConnectedRepurposer(xfer.NeuralNetworkRepurposer):\n",
    "    def __init__(self, source_model: mx.mod.Module, transfer_layer_name, num_nodes, target_class_count,\n",
    "                 context_function=mx.context.cpu, num_devices=1, batch_size=64, num_epochs=5):\n",
    "        super().__init__(source_model, context_function, num_devices, batch_size, num_epochs)\n",
    "        \n",
    "        # initialse parameters\n",
    "        self.transfer_layer_name = transfer_layer_name\n",
    "        self.num_nodes = num_nodes\n",
    "        self.target_class_count = target_class_count\n",
    "    \n",
    "    def _get_target_symbol(self, source_model_layer_names):\n",
    "        # Check if 'transfer_layer_name' is present in source model\n",
    "        if self.transfer_layer_name not in source_model_layer_names:\n",
    "            raise ValueError('transfer_layer_name: {} not found in source model'.format(self.transfer_layer_name))\n",
    "\n",
    "        # Create target symbol by transferring layers from source model up to 'transfer_layer_name'\n",
    "        transfer_layer_key = self.transfer_layer_name + '_output'  # layer key with output suffix to lookup mxnet symbol group\n",
    "        source_symbol = self.source_model.symbol.get_internals()\n",
    "        target_symbol = source_symbol[transfer_layer_key]\n",
    "        return target_symbol\n",
    "    \n",
    "    # All Neural Network Repurposers must implement this function which takes a training iterator and returns an MXNet Module\n",
    "    def _create_target_module(self, train_iterator: mx.io.DataIter):\n",
    "        # Create model handler to manipulate the source model\n",
    "        model_handler = xfer.model_handler.ModelHandler(self.source_model, self.context_function, self.num_devices)\n",
    "\n",
    "        # Create target symbol by transferring layers from source model up to and including 'transfer_layer_name'\n",
    "        target_symbol = self._get_target_symbol(model_handler.layer_names)\n",
    "\n",
    "        # Update model handler by replacing source symbol with target symbol\n",
    "        # and cleaning up weights of layers that were not transferred\n",
    "        model_handler.update_sym(target_symbol)\n",
    "\n",
    "        # Add a fully connected layer (with nodes equal to number of target classes) and a softmax output layer on top\n",
    "        fully_connected_layer1 = mx.sym.FullyConnected(num_hidden=self.num_nodes, name='fc_rep')\n",
    "        fully_connected_layer2 = mx.sym.FullyConnected(num_hidden=self.target_class_count, name='fc_from_fine_tune_repurposer')\n",
    "        softmax_output_layer = mx.sym.SoftmaxOutput(name=train_iterator.provide_label[0][0].replace('_label', ''))\n",
    "        model_handler.add_layer_top([fully_connected_layer1, fully_connected_layer2,  softmax_output_layer])\n",
    "\n",
    "        # Get fixed layers\n",
    "        conv_layer_names = model_handler.get_layer_names_matching_type('Convolution')\n",
    "        conv_layer_params = model_handler.get_layer_parameters(conv_layer_names)\n",
    "        \n",
    "        # Create and return target mxnet module using the new symbol and params\n",
    "        return model_handler.get_module(train_iterator, fixed_layer_parameters=conv_layer_params)\n",
    "    \n",
    "    # To be serialisable, Neural Network Repurposers require get_params, get_attributes, set_attributes as shown above"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# instantiate repurposer\n",
    "repurposer2Fc = Add2FullyConnectedRepurposer(source_model, transfer_layer_name='fc7', num_nodes=64, target_class_count=num_classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_iterator.reset()\n",
    "repurposer2Fc.repurpose(train_iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "results = repurposer2Fc.predict_label(test_iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "             precision    recall  f1-score   support\n",
      "\n",
      "          0       1.00      0.50      0.67         2\n",
      "          1       1.00      1.00      1.00         2\n",
      "          2       1.00      1.00      1.00         2\n",
      "          3       0.67      1.00      0.80         2\n",
      "\n",
      "avg / total       0.92      0.88      0.87         8\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(classification_report(y_pred=results, y_true=test_labels))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [default]",
   "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.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}