mnist semisup modularized.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import argparse\n",
    "import logging\n",
    "logging.basicConfig(level=logging.ERROR)\n",
    "from common import find_mxnet, data, fit\n",
    "from common.util import download_file\n",
    "import mxnet as mx\n",
    "import numpy as np\n",
    "import gzip, struct\n",
    "import time\n",
    "from mxnet.symbol import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser(description=\"train mnist\",\n",
    "                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n",
    "parser.add_argument('--num-classes', type=int, default=10,\n",
    "                        help='the number of classes')\n",
    "fit.add_fit_args(parser)\n",
    "parser.set_defaults(\n",
    "    # train\n",
    "    gpus           = '6',\n",
    "    batch_size     = 100,\n",
    "    disp_batches   = 100,\n",
    "    num_epochs     = 20000,\n",
    "    num_examples   = 100, \n",
    "    wd             = 1e-4,\n",
    "    lr             = .001,\n",
    "    lr_factor      = .33,\n",
    "    optimizer      = 'adam',\n",
    "    lr_step_epochs = '5000,10000,15000',\n",
    ")\n",
    "args = parser.parse_args(\"\")\n",
    "unsup_multiplier = 1\n",
    "labeled_per_class = 10\n",
    "sample_seed = 47\n",
    "val_interval = 100\n",
    "num_unsup_examples = 60000\n",
    "num_sup_examples = labeled_per_class * args.num_classes\n",
    "sup_batch_size = args.batch_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[5840 5656  576 3662 3626 4448 4352 5659 4106 4176]\n",
      "[1103 5635 3920 6070 6075 3763 6509 1527 1191 2650]\n",
      "[5128 5073 2352 2550  164 5421 5858 2597 1559 3745]\n",
      "[1750 1338 2470 1011 2191  774 4962 5773 6060 1810]\n",
      "[2295 5622 5409  555 5314 2464  909 3597 1125 3279]\n",
      "[5183 2728 2781 3977  345 4871 4113 1403 2388 1468]\n",
      "[5142 1629  953 4917  950  672 2646 5433 3551 5325]\n",
      "[5726 2479 1952 6120 4552 2228 2253 4124 4963 3627]\n",
      "[4854  222  260 4818 1262 3203 5607  990 5167 1794]\n",
      "[3534   41 4560 5218 4152  972 5041 2605  640 3111]\n"
     ]
    }
   ],
   "source": [
    "def read_data(label, image):\n",
    "    \"\"\"\n",
    "    download and read data into numpy\n",
    "    \"\"\"\n",
    "    base_url = 'http://yann.lecun.com/exdb/mnist/'\n",
    "    with gzip.open(download_file(base_url+label, os.path.join('data',label))) as flbl:\n",
    "        magic, num = struct.unpack(\">II\", flbl.read(8))\n",
    "        label = np.fromstring(flbl.read(), dtype=np.int8)\n",
    "    with gzip.open(download_file(base_url+image, os.path.join('data',image)), 'rb') as fimg:\n",
    "        magic, num, rows, cols = struct.unpack(\">IIII\", fimg.read(16))\n",
    "        image = np.fromstring(fimg.read(), dtype=np.uint8).reshape(len(label), rows, cols)\n",
    "    return (label, image)\n",
    "\n",
    "\n",
    "def to4d(img):\n",
    "    \"\"\"\n",
    "    reshape to 4D arrays\n",
    "    \"\"\"\n",
    "    return img.reshape(img.shape[0], 1, 28, 28).astype(np.float32)/255\n",
    "\n",
    "def sample_by_label(images, labels, n_per_label, num_labels, seed=None):\n",
    "    \"\"\"Extract equal number of sampels per class.\"\"\"\n",
    "    res_img = []\n",
    "    res_lbl = []\n",
    "    rng = np.random.RandomState(seed=seed)\n",
    "    for i in range(num_labels):\n",
    "        a = images[labels == i]\n",
    "        \n",
    "        if n_per_label == -1:  # use all available labeled data\n",
    "            res_img.append(a)\n",
    "        else:  # use randomly chosen subset\n",
    "            choice = rng.choice(len(a), n_per_label, False)\n",
    "            print(choice)\n",
    "            r = a[choice]\n",
    "            \n",
    "            res_img.append(r)\n",
    "            \n",
    "            lbls = np.ones(n_per_label) * i\n",
    "            res_lbl.append(lbls)\n",
    "    return (res_img, res_lbl)\n",
    "\n",
    "\n",
    "\"\"\"\n",
    "create data iterator with NDArrayIter\n",
    "\"\"\"\n",
    "(train_lbl, train_img) = read_data(\n",
    "        'train-labels-idx1-ubyte.gz', 'train-images-idx3-ubyte.gz')\n",
    "(val_lbl, val_img) = read_data(\n",
    "        't10k-labels-idx1-ubyte.gz', 't10k-images-idx3-ubyte.gz')\n",
    "\n",
    "(imgs, lbls) = sample_by_label(to4d(train_img), train_lbl, labeled_per_class, 10, seed = sample_seed)\n",
    "imgs = np.vstack(imgs)\n",
    "lbls = np.hstack(lbls)\n",
    "\n",
    "train_sup = mx.io.NDArrayIter(\n",
    "    imgs, lbls, args.batch_size, shuffle=True, data_name='dataSup', label_name='labelSup')\n",
    "train_unsup = mx.io.NDArrayIter(\n",
    "    to4d(train_img), label=None, batch_size=args.batch_size, \n",
    "    shuffle=True, data_name='dataUnsup')\n",
    "val = mx.io.NDArrayIter(\n",
    "    to4d(val_img), val_lbl, args.batch_size*1) # use larger test batch size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "600  times more unsup data than sup data\n"
     ]
    }
   ],
   "source": [
    "from common.multi_iterator import Multi_iterator\n",
    "    \n",
    "train = Multi_iterator(train_sup, train_unsup, unsup_multiplier, num_unsup_examples, num_sup_examples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_embedding_shapes(unsup_multiplier):\n",
    "    embedding_shapes = [mx.io.DataDesc(\"embeddings_sup\", (100,128), np.float32, \"NCHW\")]\n",
    "    for i in range(unsup_multiplier):\n",
    "        embedding_shapes = embedding_shapes + [mx.io.DataDesc(\"embeddings_unsup\"+str(i), (100,128), np.float32, \"NCHW\")]\n",
    "    return embedding_shapes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[5000.0, 10000.0, 15000.0]\n"
     ]
    }
   ],
   "source": [
    "from symbols import mnist\n",
    "from common.lba import compute_semisup_loss, logit_loss\n",
    "\n",
    "# kvstore\n",
    "kv = mx.kvstore.create(args.kv_store)\n",
    "\n",
    "t_nb = sup_batch_size * unsup_multiplier\n",
    "initializer = mx.init.Xavier(rnd_type='gaussian', factor_type=\"avg\", magnitude=2.34)\n",
    "\n",
    "lr, lr_scheduler = fit._get_lr_scheduler(args, kv)\n",
    "\n",
    "def buildEmbeddingModule():\n",
    "    dataSup = mx.symbol.Variable(name=\"dataSup\")\n",
    "    dataUnsup = []\n",
    "    for i in range(unsup_multiplier):\n",
    "        dataUnsup.append(Variable(name=\"dataUnsup\"+str(i)))\n",
    "\n",
    "    # concat data, feed both through the network\n",
    "    # then split it up again\n",
    "    data = concat(dataSup, *dataUnsup, dim=0)\n",
    "\n",
    "    embeddings = mnist.build_embeddings(data)\n",
    "    splitted = split(embeddings, num_outputs=(unsup_multiplier+1), axis=0, name='split')\n",
    "    \n",
    "    # devices for training\n",
    "    devs = mx.cpu() if args.gpus is None or args.gpus is '' else [\n",
    "        mx.gpu(int(i)) for i in args.gpus.split(',')]\n",
    "\n",
    "    data_names = ['dataSup'] + ['dataUnsup'+str(i) for i in range(unsup_multiplier)]\n",
    "        \n",
    "    # create model\n",
    "    model = mx.mod.Module(\n",
    "        context       = devs,\n",
    "        symbol        = splitted,\n",
    "        data_names    = data_names,\n",
    "        label_names   = None)\n",
    "    \n",
    "    model.bind(data_shapes=train.provide_data)         \n",
    "    model.init_params(initializer)    \n",
    "    model.init_optimizer(optimizer='adam', optimizer_params=(\n",
    "        ('learning_rate', lr), \n",
    "        ('wd', 1e-4),\n",
    "        ('rescale_grad', 0.1),\n",
    "        ('lr_scheduler', lr_scheduler)))\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "def buildLossModule():\n",
    "    supEmbeddings = Variable(name=\"embeddings_sup\")\n",
    "    labelSup = mx.symbol.Variable(name='labelSup')\n",
    "    overall_loss = []\n",
    "    \n",
    "    if unsup_multiplier >= 1:\n",
    "        unsupEmbeddings = []\n",
    "        for i in range(unsup_multiplier):\n",
    "            unsupEmbeddings.append(Variable(name=\"embeddings_unsup\"+str(i)))\n",
    "    \n",
    "        unsupEmbeddings = concat(*unsupEmbeddings, dim=0)\n",
    "            \n",
    "        (walker_loss, visit_loss) = compute_semisup_loss(supEmbeddings, unsupEmbeddings, labelSup, t_nb, \n",
    "                                                     walker_weight=1.0, visit_weight=1.0)\n",
    "        overall_loss = [walker_loss, visit_loss]\n",
    "        \n",
    "    overall_loss = [logit_loss(supEmbeddings, labelSup, 10)] + overall_loss\n",
    "    \n",
    "    # todo maybe use gpu\n",
    "    devs = mx.cpu()\n",
    "\n",
    "    # create module\n",
    "    model = mx.mod.Module(\n",
    "        context = devs,\n",
    "        symbol  = Group(overall_loss),\n",
    "        data_names = ['embeddings_sup']+['embeddings_unsup'+str(i) for i in range(unsup_multiplier)],\n",
    "        label_names = ['labelSup'])\n",
    "    \n",
    "    # allocate memory by given the input data and label shapes\n",
    "    model.bind(data_shapes=get_embedding_shapes(unsup_multiplier), label_shapes=train.provide_label,\n",
    "            inputs_need_grad=True)\n",
    "    \n",
    "    model.init_params(initializer)  \n",
    "    model.init_optimizer(optimizer='adam', optimizer_params=(\n",
    "        ('learning_rate', lr), \n",
    "        ('rescale_grad', 0.1),\n",
    "        ('wd', 1e-4),\n",
    "        ('lr_scheduler', lr_scheduler)))\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WrapperModule(mx.mod.BaseModule):\n",
    "    \n",
    "    def __init__(self, embedding_module_builder, loss_module_builder, logger=None):\n",
    "        super(mx.mod.BaseModule, self).__init__()\n",
    "        \n",
    "        self.logger = logging\n",
    "        self._symbol = None\n",
    "        self.binded = True\n",
    "        self.params_initialized = True\n",
    "        \n",
    "        self.embedding_module = embedding_module_builder()\n",
    "        self.loss_module = loss_module_builder()\n",
    "        \n",
    "    def forward(self, data_batch, is_train=True):\n",
    "        self.embedding_module.forward(data_batch, is_train=is_train)\n",
    "        embeddings = self.embedding_module.get_outputs()\n",
    "\n",
    "        self.loss_module.forward(mx.io.DataBatch(embeddings, label=data_batch.label), is_train=is_train)\n",
    "\n",
    "        \n",
    "    def backward(self):\n",
    "        self.loss_module.backward()\n",
    "        grads = self.loss_module.get_input_grads()\n",
    "\n",
    "        self.embedding_module.backward(grads)\n",
    "        \n",
    "    def update(self):\n",
    "        self.embedding_module.update()\n",
    "        self.loss_module.update()\n",
    "        \n",
    "    def bind(self, *args, **kwargs):\n",
    "        1\n",
    "        # do nothing\n",
    "        \n",
    "    def init_params(self, *args, **kwargs):\n",
    "        1\n",
    "        # do nothing\n",
    "        \n",
    "    def get_params(self, *args, **kwargs):\n",
    "        return (1,2)\n",
    "        # do nothing\n",
    "        \n",
    "    def set_params(self, *args, **kwargs):\n",
    "        1\n",
    "        # do nothing\n",
    "    def init_optimizer(self, *args, **kwargs):\n",
    "        1\n",
    "        # do nothing\n",
    "    def update_metric(self, *args, **kwargs):\n",
    "        1\n",
    "        # todo implement\n",
    "        \n",
    "    def get_outputs(self):\n",
    "        return self.loss_module.get_outputs()\n",
    "        \n",
    "    def score(self, eval_data, eval_metric, num_batch=None, batch_end_callback=None,\n",
    "              score_end_callback=None,\n",
    "              reset=True, epoch=0):\n",
    "        \n",
    "        eval_data.reset()\n",
    "        \n",
    "        pred = self.predict(eval_data)\n",
    "        res = pred[0].asnumpy().argmax(axis=1)\n",
    "        \n",
    "        acc = (res == val_lbl).mean()  # todo get labels from batch\n",
    "        \n",
    "        return acc\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_metrics = Multi_Accuracy(num= 3 if unsup_multiplier >= 1 else 1)\n",
    "                    \n",
    "def fit_model(args, data, **kwargs):\n",
    "    \"\"\"\n",
    "    train a model\n",
    "    args : argparse returns\n",
    "    data_loader : function that returns the train and val data iterators\n",
    "    \"\"\"\n",
    "    # logging\n",
    "    head = '%(asctime)-15s Node[' + str(kv.rank) + '] %(message)s'\n",
    "    logging.basicConfig(level=logging.DEBUG, format=head)\n",
    "    logging.info('start with arguments %s', args)\n",
    "    batch_size = sup_batch_size * (unsup_multiplier + 1)\n",
    "    \n",
    "    # data iterators\n",
    "    (train, val) = data\n",
    "\n",
    "    # load model\n",
    "    if 'arg_params' in kwargs and 'aux_params' in kwargs:\n",
    "        arg_params = kwargs['arg_params']\n",
    "        aux_params = kwargs['aux_params']\n",
    "    else: \n",
    "        sym, arg_params, aux_params = fit._load_model(args, kv.rank)\n",
    "        \n",
    "    # save model\n",
    "    checkpoint = fit._save_model(args, kv.rank)\n",
    "\n",
    "    model = WrapperModule(buildEmbeddingModule, buildLossModule)\n",
    "    \n",
    "    def validate_model(epoch, *args):\n",
    "        if epoch % val_interval != 0: \n",
    "            return\n",
    "        res = model.score(val, eval_metrics)\n",
    "        #TODO: pull this into default\n",
    "        print('Epoch[%d] Validation-accuracy=%f' % (epoch,  res))\n",
    "\n",
    "    #monitor = mx.mon.Monitor(interval=1000, pattern='.*aba_backward.*') \n",
    "    monitor = mx.mon.Monitor(interval=1000, pattern='.*') \n",
    "    \n",
    "    # callbacks that run after each batch\n",
    "    batch_end_callbacks = [mx.callback.Speedometer(batch_size, args.disp_batches)]\n",
    "    \n",
    "    epoch_end_callbacks = validate_model\n",
    "\n",
    "    # run\n",
    "    model.fit(train,\n",
    "        begin_epoch        = args.load_epoch if args.load_epoch else 0,\n",
    "        num_epoch          = args.num_epochs,\n",
    "        #eval_data          = val,\n",
    "        eval_metric        = eval_metrics,\n",
    "        kvstore            = kv,\n",
    "        arg_params         = arg_params,\n",
    "        aux_params         = aux_params,\n",
    "        batch_end_callback = batch_end_callbacks,\n",
    "        epoch_end_callback = epoch_end_callbacks,\n",
    "        allow_missing      = True\n",
    "        #monitor            = monitor\n",
    "             )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch[0] Validation-accuracy=0.102700\n",
      "Epoch[100] Validation-accuracy=0.925400\n",
      "Epoch[200] Validation-accuracy=0.956000\n",
      "Epoch[300] Validation-accuracy=0.962900\n",
      "Epoch[400] Validation-accuracy=0.966900\n",
      "Epoch[500] Validation-accuracy=0.971600\n",
      "Epoch[600] Validation-accuracy=0.973500\n",
      "Epoch[700] Validation-accuracy=0.974400\n"
     ]
    }
   ],
   "source": [
    "train.reset()\n",
    "fit_model(args, (train, val))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'mxnet.metric' from '/usr/local/lib/python3.4/dist-packages/mxnet/metric.py'>"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class Multi_Accuracy(mx.metric.EvalMetric):\n",
    "    \"\"\"Calculate accuracies of multi label\"\"\"\n",
    "\n",
    "    def __init__(self, num=None):\n",
    "        super(Multi_Accuracy, self).__init__('multi-accuracy', num)\n",
    "\n",
    "    def update(self, labels, preds):\n",
    "\n",
    "        #for i in range(len(preds)):\n",
    "        for i in range(1):\n",
    "            pred_label = mx.nd.argmax_channel(preds[i]).asnumpy().astype('int32')\n",
    "            label = labels[0].asnumpy().astype('int32')\n",
    "\n",
    "            #mx.metric.check_label_shapes(label, pred_label)\n",
    "            \n",
    "            #print((pred_label.flat == label.flat).sum())\n",
    "            #print(len(pred_label.flat))\n",
    "\n",
    "            \n",
    "            self.sum_metric[i] += (pred_label.flat == label.flat).sum()\n",
    "            self.num_inst[i] += len(pred_label.flat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print(labels.asnumpy())\n",
    "model.get_outputs()[1].asnumpy()[0:10,0:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}