transfer_learning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install gluoncv # -i https://opentuna.cn/pypi/web/simple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. Transfer Learning with Your Own Image Dataset\n",
    "=======================================================\n",
    "\n",
    "Dataset size is a big factor in the performance of deep learning models.\n",
    "``ImageNet`` has over one million labeled images, but\n",
    "we often don't have so much labeled data in other domains.\n",
    "Training a deep learning models on small datasets may lead to severe overfitting.\n",
    "\n",
    "Transfer learning is a technique that addresses this problem.\n",
    "The idea is simple: we can start training with a pre-trained model,\n",
    "instead of starting from scratch.\n",
    "As Isaac Newton said, \"If I have seen further it is by standing on the\n",
    "shoulders of Giants\".\n",
    "\n",
    "In this tutorial, we will explain the basics of transfer\n",
    "learning, and apply it to the ``MINC-2500`` dataset.\n",
    "\n",
    "Data Preparation\n",
    "----------------\n",
    "\n",
    "`MINC <http://opensurfaces.cs.cornell.edu/publications/minc/>`__ is\n",
    "short for Materials in Context Database, provided by Cornell.\n",
    "``MINC-2500`` is a resized subset of ``MINC`` with 23 classes, and 2500\n",
    "images in each class. It is well labeled and has a moderate size thus is\n",
    "perfect to be our example.\n",
    "\n",
    "|image-minc|\n",
    "\n",
    "To start, we first download ``MINC-2500`` from\n",
    "`here <http://opensurfaces.cs.cornell.edu/publications/minc/>`__.\n",
    "Suppose we have the data downloaded to ``~/data/`` and\n",
    "extracted to ``~/data/minc-2500``.\n",
    "\n",
    "After extraction, it occupies around 2.6GB disk space with the following\n",
    "structure:\n",
    "\n",
    "::\n",
    "\n",
    "    minc-2500\n",
    "    ├── README.txt\n",
    "    ├── categories.txt\n",
    "    ├── images\n",
    "    └── labels\n",
    "\n",
    "The ``images`` folder has 23 sub-folders for 23 classes, and ``labels``\n",
    "folder contains five different splits for training, validation, and test.\n",
    "\n",
    "We have written a script to prepare the data for you:\n",
    "\n",
    ":download:`Download prepare_minc.py<../../../scripts/classification/finetune/prepare_minc.py>`\n",
    "\n",
    "Run it with\n",
    "\n",
    "::\n",
    "\n",
    "    python prepare_minc.py --data ~/data/minc-2500 --split 1\n",
    "\n",
    "Now we have the following structure:\n",
    "\n",
    "::\n",
    "\n",
    "    minc-2500\n",
    "    ├── categories.txt\n",
    "    ├── images\n",
    "    ├── labels\n",
    "    ├── README.txt\n",
    "    ├── test\n",
    "    ├── train\n",
    "    └── val\n",
    "\n",
    "In order to go through this tutorial within a reasonable amount of time,\n",
    "we have prepared a small subset of the ``MINC-2500`` dataset,\n",
    "but you should substitute it with the original dataset for your experiments.\n",
    "We can download and extract it with:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hyperparameters\n",
    "----------\n",
    "\n",
    "First, let's import all other necessary libraries.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import mxnet as mx\n",
    "import numpy as np\n",
    "import os, time, shutil\n",
    "\n",
    "from mxnet import gluon, image, init, nd\n",
    "from mxnet import autograd as ag\n",
    "from mxnet.gluon import nn\n",
    "from mxnet.gluon.data.vision import transforms\n",
    "from gluoncv.utils import makedirs\n",
    "from gluoncv.model_zoo import get_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We set the hyperparameters as following:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "classes = 5\n",
    "\n",
    "epochs = 100\n",
    "lr = 0.001\n",
    "per_device_batch_size = 32\n",
    "momentum = 0.9\n",
    "wd = 0.0001\n",
    "\n",
    "lr_factor = 0.75\n",
    "lr_steps = [10, 20, 30, np.inf]\n",
    "\n",
    "num_gpus = 1\n",
    "num_workers = 8\n",
    "ctx = [mx.gpu(i) for i in range(num_gpus)] if num_gpus > 0 else [mx.cpu()]\n",
    "batch_size = per_device_batch_size * max(num_gpus, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Things to keep in mind:\n",
    "\n",
    "1. ``epochs = 5`` is just for this tutorial with the tiny dataset. please change it to a larger number in your experiments, for instance 40.\n",
    "2. ``per_device_batch_size`` is also set to a small number. In your experiments you can try larger number like 64.\n",
    "3. remember to tune ``num_gpus`` and ``num_workers`` according to your machine.\n",
    "4. A pre-trained model is already in a pretty good status. So we can start with a small ``lr``.\n",
    "\n",
    "Data Augmentation\n",
    "-----------------\n",
    "\n",
    "In transfer learning, data augmentation can also help.\n",
    "We use the following augmentation in training:\n",
    "\n",
    "2. Randomly crop the image and resize it to 224x224\n",
    "3. Randomly flip the image horizontally\n",
    "4. Randomly jitter color and add noise\n",
    "5. Transpose the data from height*width*num_channels to num_channels*height*width, and map values from [0, 255] to [0, 1]\n",
    "6. Normalize with the mean and standard deviation from the ImageNet dataset.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "jitter_param = 0.4\n",
    "lighting_param = 0.1\n",
    "\n",
    "transform_train = transforms.Compose([\n",
    "    transforms.RandomResizedCrop(224),\n",
    "    transforms.RandomFlipLeftRight(),\n",
    "    transforms.RandomColorJitter(brightness=jitter_param, contrast=jitter_param,\n",
    "                                 saturation=jitter_param),\n",
    "    transforms.RandomLighting(lighting_param),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.Resize(256),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the data augmentation functions, we can define our data loaders:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -O image_classification.zip \"https://datalab.s3.amazonaws.com/data/image_classification.zip?AWSAccessKeyId=AKIAYNUCDPLSDWHHQJ7Y&Signature=PTOluDRsNdp2riiej6Q9DagCAiw%3D&Expires=1648453417\"\n",
    "!unzip -q image_classification.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p data/train\n",
    "!mkdir -p data/test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "base_dir = 'image_classification'\n",
    "filenames = os.listdir(base_dir)\n",
    "class_names = []\n",
    "for filename in filenames:\n",
    "    if os.path.isdir(os.path.join(base_dir, filename)) and not filename.startswith('.'):\n",
    "        class_names.append(filename)\n",
    "        if not os.path.exists(os.path.join('data/train/', filename)):\n",
    "            os.mkdir(os.path.join('data/train/', filename))\n",
    "            os.mkdir(os.path.join('data/test/', filename))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "for name in class_names:\n",
    "    filenames = os.listdir(os.path.join(base_dir, name))\n",
    "    print(name, len(filenames))\n",
    "    train_filenames, test_filenames = train_test_split(filenames, test_size=0.3)\n",
    "    for filename in train_filenames:\n",
    "        os.system('cp '+os.path.join(base_dir, name, filename)+' '+os.path.join('data/train/', name, filename))\n",
    "    for filename in test_filenames:\n",
    "        os.system('cp '+os.path.join(base_dir, name, filename)+' '+os.path.join('data/test/', name, filename))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "path = './data'\n",
    "train_path = os.path.join(path, 'train')\n",
    "val_path = os.path.join(path, 'test')\n",
    "test_path = os.path.join(path, 'test')\n",
    "\n",
    "train_data = gluon.data.DataLoader(\n",
    "    gluon.data.vision.ImageFolderDataset(train_path).transform_first(transform_train),\n",
    "    batch_size=batch_size, shuffle=True, num_workers=num_workers)\n",
    "\n",
    "val_data = gluon.data.DataLoader(\n",
    "    gluon.data.vision.ImageFolderDataset(val_path).transform_first(transform_test),\n",
    "    batch_size=batch_size, shuffle=False, num_workers = num_workers)\n",
    "\n",
    "test_data = gluon.data.DataLoader(\n",
    "    gluon.data.vision.ImageFolderDataset(test_path).transform_first(transform_test),\n",
    "    batch_size=batch_size, shuffle=False, num_workers = num_workers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(gluon.data.vision.ImageFolderDataset(train_path).synsets)\n",
    "print(gluon.data.vision.ImageFolderDataset(val_path).synsets)\n",
    "print(gluon.data.vision.ImageFolderDataset(test_path).synsets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that only ``train_data`` uses ``transform_train``, while\n",
    "``val_data`` and ``test_data`` use ``transform_test`` to produce deterministic\n",
    "results for evaluation.\n",
    "\n",
    "Model and Trainer\n",
    "-----------------\n",
    "\n",
    "We use a pre-trained ``ResNet50_v2`` model, which has balanced accuracy and\n",
    "computation cost.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "model_name = 'ResNet50_v2'\n",
    "# model_name = 'ResNet152_v1d'\n",
    "finetune_net = get_model(model_name, pretrained=True)\n",
    "with finetune_net.name_scope():\n",
    "    finetune_net.output = nn.Dense(classes)\n",
    "finetune_net.output.initialize(init.Xavier(), ctx = ctx)\n",
    "finetune_net.collect_params().reset_ctx(ctx)\n",
    "finetune_net.hybridize()\n",
    "\n",
    "trainer = gluon.Trainer(finetune_net.collect_params(), 'sgd', {\n",
    "                        'learning_rate': lr, 'momentum': momentum, 'wd': wd})\n",
    "metric = mx.metric.Accuracy()\n",
    "L = gluon.loss.SoftmaxCrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an illustration of the pre-trained model\n",
    "and our newly defined model:\n",
    "\n",
    "|image-model|\n",
    "\n",
    "Specifically, we define the new model by::\n",
    "\n",
    "1. load the pre-trained model\n",
    "2. re-define the output layer for the new task\n",
    "3. train the network\n",
    "\n",
    "This is called \"fine-tuning\", i.e. we have a model trained on another task,\n",
    "and we would like to tune it for the dataset we have in hand.\n",
    "\n",
    "We define a evaluation function for validation and testing.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def test(net, val_data, ctx):\n",
    "    metric = mx.metric.Accuracy()\n",
    "    for i, batch in enumerate(val_data):\n",
    "        data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0, even_split=False)\n",
    "        label = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0, even_split=False)\n",
    "        outputs = [net(X) for X in data]\n",
    "        metric.update(label, outputs)\n",
    "        # print(label, outputs)\n",
    "\n",
    "    return metric.get()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training Loop\n",
    "-------------\n",
    "\n",
    "Following is the main training loop. It is the same as the loop in\n",
    "`CIFAR10 <dive_deep_cifar10.html>`__\n",
    "and ImageNet.\n",
    "\n",
    "<div class=\"alert alert-info\"><h4>Note</h4><p>Once again, in order to go through the tutorial faster, we are training on a small\n",
    "    subset of the original ``MINC-2500`` dataset, and for only 5 epochs. By training on the\n",
    "    full dataset with 40 epochs, it is expected to get accuracy around 80% on test data.</p></div>\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "lr_counter = 0\n",
    "num_batch = len(train_data)\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    if epoch == lr_steps[lr_counter]:\n",
    "        trainer.set_learning_rate(trainer.learning_rate*lr_factor)\n",
    "        lr_counter += 1\n",
    "\n",
    "    tic = time.time()\n",
    "    train_loss = 0\n",
    "    metric.reset()\n",
    "\n",
    "    for i, batch in enumerate(train_data):\n",
    "#         print(i)\n",
    "        data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0, even_split=False)\n",
    "        label = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0, even_split=False)\n",
    "#         print(label)\n",
    "        with ag.record():\n",
    "            outputs = [finetune_net(X) for X in data]\n",
    "            loss = [L(yhat, y) for yhat, y in zip(outputs, label)]\n",
    "        for l in loss:\n",
    "            l.backward()\n",
    "\n",
    "        trainer.step(batch_size)\n",
    "        train_loss += sum([l.mean().asscalar() for l in loss]) / len(loss)\n",
    "\n",
    "        metric.update(label, outputs)\n",
    "\n",
    "    _, train_acc = metric.get()\n",
    "    train_loss /= num_batch\n",
    "\n",
    "    _, val_acc = test(finetune_net, val_data, ctx)\n",
    "\n",
    "    print('[Epoch %d] Train-acc: %.3f, loss: %.3f | Val-acc: %.3f | time: %.1f' %\n",
    "             (epoch, train_acc, train_loss, val_acc, time.time() - tic))\n",
    "\n",
    "_, test_acc = test(finetune_net, test_data, ctx)\n",
    "print('[Finished] Test-acc: %.3f' % (test_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('[Finished] Test-acc: %.3f' % (test_acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir endpoint/model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "finetune_net.save_parameters('endpoint/model/model-0000.params')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next\n",
    "----\n",
    "\n",
    "Now that you have learned to muster the power of transfer\n",
    "learning, to learn more about training a model on\n",
    "ImageNet, please read `this tutorial <dive_deep_imagenet.html>`__.\n",
    "\n",
    "The idea of transfer learning is the basis of\n",
    "`object detection <../examples_detection/index.html>`_ and\n",
    "`semantic segmentation <../examples_segmentation/index.html>`_,\n",
    "the next two chapters of our tutorial.\n",
    "\n",
    ".. |image-minc| image:: https://raw.githubusercontent.com/dmlc/web-data/master/gluoncv/datasets/MINC-2500.png\n",
    ".. |image-model| image:: https://zh.gluon.ai/_images/fine-tuning.svg\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_mxnet_p36",
   "language": "python",
   "name": "conda_mxnet_p36"
  },
  "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}