Example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "85d3232e",
   "metadata": {},
   "source": [
    "## Branching Point Detector\n",
    "Given a backbone model specified in the format of ```*.prototxt```, the branching point detector will automatically divide it into sequential blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28a7fd07",
   "metadata": {},
   "outputs": [],
   "source": [
    "from main.auto_models import MTSeqBackbone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9e59364",
   "metadata": {},
   "outputs": [],
   "source": [
    "prototxt = 'models/deeplab_resnet34_adashare.prototxt' # MobileNetV2: models/mobilenetv2.prototxt\n",
    "backbone = MTSeqBackbone(prototxt)\n",
    "B = len(backbone.basic_blocks)\n",
    "print('The number of blocks is {}.'.format(B))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5da4bc78",
   "metadata": {},
   "source": [
    "The user can __further specify coarse-grained branching points__ based on the auto-generated branching points by defining a mapping dictionary.\n",
    "The following example maps the 32 branching points to 5 coarser ones:\n",
    "\n",
    "| coarse |    0 |   1 |     2   |   3   |  4 |\n",
    "|:------:|:------:|:---------:|:---------:|:----:|:--:|\n",
    "|  fined | 0 | 1,2,3 | 4,5,6,7 | 8,9,10,11,12,13 | 14,15,16 |\n",
    "\n",
    "__Note:__ The ```mapping``` dict has the last element ```5:[17]``` to indicate the all-shared models. In other words, if you're mapping $M$ branching points to $N$ coarser ones, ```mapping``` will contains $N+1$ elements in which the last one is ```N:[M]```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ac41e11",
   "metadata": {},
   "outputs": [],
   "source": [
    "coarse_B = 5\n",
    "mapping = {0:[0], 1:[1,2,3], 2:[4,5,6,7], 3:[8,9,10,11,12,13], 4:[14,15,16], 5:[17]}\n",
    "# mapping = {0:[0,1,2,3,4,5,6], 1:[7,8,9,10,11,12,13,14,15,16,17], 2:[18,19,20,21,22], \n",
    "#            3:[23,24,25,26,27,28,29,30], 4:[31], 5:[32]} # mapping for MobileNetV2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38d1a6bd",
   "metadata": {},
   "source": [
    "## Design Spce Enumerator\n",
    "Given the number of tasks $T$ and the number of branching points $B$, the design space enumerator could explore the tree-structured multi-task model architectures space completely."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "750c501c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from main.algorithms import enumerator "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f8ad4b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "T = 3 # NYUv2 has 3 tasks, Taskonomy has 5 tasks\n",
    "layout_list = enumerator(T, coarse_B)\n",
    "print('There are {} layouts in the design space.'.format(len(layout_list)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9da43f0c",
   "metadata": {},
   "source": [
    "## Task Accuracy Estimator\n",
    "For each layout in the design space, we will estimate its task accuracy from the task accuacy of associated 2-task models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "059b1fcb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from main.algorithms import reorg_two_task_results, compute_weights, metric_inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65408a60",
   "metadata": {},
   "source": [
    "### Step 1: Load 2-task models results\n",
    "\n",
    "__The task accuracy of all the 2-task models should be stored in excel and organized as the following example.__\n",
    "\n",
    "* Each column represents different 2-task combinations. \n",
    "For $(a,b)-i: i \\in \\{0,1\\}$, $(a,b)$ refers to the 2-task model of task $a$ and task $b$, and $-i$ means the current column is the accuracy of $i$-th task -- $0$ is task $a$, $1$ is task $b$.\n",
    "\n",
    "* Each row represents the branching points of the 2-task models.\n",
    "Notice that $0$ means independent models, while $B$ means all-shared models.\n",
    "\n",
    "More examples can be found in the folder ```2task/*.xlsx```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ab0c728",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "two_task_pd = pd.read_excel('2task/NYUv2_2task_metrics_resnet_1129_val_acc.xlsx',engine='openpyxl',index_col=0)\n",
    "two_task_pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff33a74f",
   "metadata": {},
   "source": [
    "### Step 2: Compute score weights\n",
    "\n",
    "The 2-task results will be reorganized by ```reorg_two_task_results``` for the further computation, and the task weights for the layout scores are computed by ```compute_weights```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64d8b565",
   "metadata": {},
   "outputs": [],
   "source": [
    "two_task_metrics = reorg_two_task_results(two_task_pd, T, coarse_B)\n",
    "score_weights = compute_weights(two_task_pd, T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c70fdd3",
   "metadata": {},
   "source": [
    "### Step 3: Compute layout score from accociated 2-task models\n",
    "\n",
    "For each layout in the design space ```layout_list```, figure out the accociated 2-task models by ```metric_inference```, then set the final score by ```L.set_score_weighted```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1362a4ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run for all L\n",
    "for L in layout_list:\n",
    "    print('Layout: {}'.format(L))\n",
    "    \n",
    "    subtree = metric_inference(L, two_task_metrics)\n",
    "    print('Associated 2-task models for each task: {}'.format(subtree))\n",
    "    \n",
    "    L.set_score_weighted(score_weights)\n",
    "    print('Final Score: {:.4f}'.format(L.score))\n",
    "    \n",
    "    print('=' * 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5da94b6",
   "metadata": {},
   "source": [
    "### Step 4: Sort the layouts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2015175",
   "metadata": {},
   "outputs": [],
   "source": [
    "layout_order = sorted(range(len(layout_list)), key=lambda k: layout_list[k].score,reverse=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c361ec3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(0,len(layout_order)):\n",
    "    print('Layout Idx: {}'.format(layout_order[i]))\n",
    "    L = layout_list[layout_order[i]]\n",
    "    print('Layout: {}'.format(L))\n",
    "    print('Final Score: {:.4f}'.format(L.score))\n",
    "    print('=' * 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c85be197",
   "metadata": {},
   "source": [
    "## Appendix 1: Dataloader, Loss, and Metrics\n",
    "\n",
    "We provide dataloader, loss functions, and metrics evaluations for NYUv2 and Taskonomy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8445788f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from data.nyuv2_dataloader_adashare import NYU_v2\n",
    "from data.taskonomy_dataloader_adashare import Taskonomy\n",
    "from data.pixel2pixel_loss import NYUCriterions, TaskonomyCriterions\n",
    "from data.pixel2pixel_metrics import NYUMetrics, TaskonomyMetrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee64f416",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataroot = 'data/NYUv2' # Your root\n",
    "\n",
    "criterionDict = {}\n",
    "metricDict = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20c0da2a",
   "metadata": {},
   "source": [
    "### NYUv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3d28730",
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks = ['segment_semantic','normal','depth_zbuffer']\n",
    "cls_num = {'segment_semantic': 40, 'normal':3, 'depth_zbuffer': 1}\n",
    "\n",
    "dataset = NYU_v2(dataroot, 'train', crop_h=321, crop_w=321)\n",
    "trainDataloader = DataLoader(dataset, 16, shuffle=True)\n",
    "\n",
    "dataset = NYU_v2(dataroot, 'test', crop_h=321, crop_w=321)\n",
    "valDataloader = DataLoader(dataset, 16, shuffle=True)\n",
    "\n",
    "for task in tasks:\n",
    "    criterionDict[task] = NYUCriterions(task)\n",
    "    metricDict[task] = NYUMetrics(task)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6efd4fb",
   "metadata": {},
   "source": [
    "### Taskonomy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dab6d91",
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks = ['segment_semantic','normal','depth_zbuffer','keypoints2d','edge_texture']\n",
    "cls_num = {'segment_semantic': 17, 'normal': 3, 'depth_zbuffer': 1, 'keypoints2d': 1, 'edge_texture': 1}\n",
    "    \n",
    "dataset = Taskonomy(dataroot, 'train', crop_h=224, crop_w=224)\n",
    "trainDataloader = DataLoader(dataset, batch_size=16, shuffle=True)\n",
    "\n",
    "dataset = Taskonomy(dataroot, 'test_small', crop_h=224, crop_w=224)\n",
    "valDataloader = DataLoader(dataset, batch_size=16, shuffle=True)\n",
    "\n",
    "for task in tasks:\n",
    "    criterionDict[task] = TaskonomyCriterions(task, dataroot)\n",
    "    metricDict[task] = TaskonomyMetrics(task, dataroot)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68c60ef8",
   "metadata": {},
   "source": [
    "## Appendix 2: 2-Task Models & N-Task Models\n",
    "We need to train all the 2-task models at different branching points for the performance table, and the n-task models we select after estimating and sorting their task accuracy. Therefore we also provide __a model generator that can automatically build up the 2-task models based on the given branching points, and the n-task models based on the given layout__."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "528836f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from main.auto_models import MTSeqModel\n",
    "from main.algorithms import coarse_to_fined"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca3d8310",
   "metadata": {},
   "source": [
    "### 2-Task Model\n",
    "* Inputs: prototxt, the branching point, the number of branching points, the feature dimension and the number of class for task heads\n",
    "* __Note:__\n",
    "    * Given a coarse branch point, we can convert it to a fined branch point from the mapping.\n",
    "    * The feature dimension can be defined by the user or derived from the backbone model.\n",
    "    * Remember to select 2 tasks from the task set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbb31a91",
   "metadata": {},
   "outputs": [],
   "source": [
    "coarse_branch = 3\n",
    "fined_branch = mapping[coarse_branch][0]\n",
    "feature_dim = backbone(torch.rand(1,3,224,224)).shape[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31d3bee7",
   "metadata": {},
   "outputs": [],
   "source": [
    "two_task = ['segment_semantic','normal'] # Select two tasks as you like\n",
    "two_cls_num = {task: cls_num[task] for task in two_task}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f17f915f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MTSeqModel(prototxt, branch=fined_branch, fined_B=B, feature_dim=feature_dim, cls_num=two_cls_num)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b5682c3",
   "metadata": {},
   "source": [
    "### N-Task Model\n",
    "* Inputs: prototxt, a layout, the feature dimension and the number of class for task heads\n",
    "* __Note:__\n",
    "    * Given a layout enumerated under the coarse branching points, we can use ```coarse_to_fined``` to convert it to a layout under the fined branching points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "134f939c",
   "metadata": {},
   "outputs": [],
   "source": [
    "coarse_layout = layout_list[45]\n",
    "fined_layout = coarse_to_fined(coarse_layout, B, mapping)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d654734a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MTSeqModel(prototxt, layout=fined_layout, feature_dim=feature_dim, cls_num=cls_num)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1670e387",
   "metadata": {},
   "source": [
    "## Appendix 3: Trainer Functions\n",
    "\n",
    "We further provide trainer functions to train the 2-task and n-task models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f1bdef5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from main.trainer import Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f6bec3e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(model.cuda(), tasks, trainDataloader, valDataloader, criterionDict, metricDict)\n",
    "trainer.train(20000)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (.conda-multitask)",
   "language": "python",
   "name": "conda-env-.conda-multitask-py"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}