tutorial-and-how-to/03-autotvm-tune-with-templates.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AutoTVM: Tune with Templates\n",
    "\n",
    "官网 doc: https://tvm.apache.org/docs/how_to/tune_with_autotvm/index.html\n",
    "\n",
    "要点:\n",
    "1. 使用 autotvm 模块优化 AlexNet 模型。与 [TVM Ansor: Template-Free Auto-Tuning](https://github.com/JackonYang/hands-on-tvm/blob/main/tutorial-and-how-to/02-tvm-ansor-template-free-tune.ipynb) 用的模型相同.\n",
    "2. autotvm 的 tune 分 2 步：(1) tune kernel; (2) tune graph.\n",
    "3. 总计 tune 约 13min，模型加速 ~5 倍。硬件: Intel(R) Xeon(R) Gold 5320 CPU。\n",
    "\n",
    "备注:\n",
    "\n",
    "1. 虽然，原理上需要写 template，但 CPU/GPU 上，使用默认 template，无脑暴力搜，结果也不错。\n",
    "2. 搜 kernel 的 tuner_obj，& 搜 graph 的 Tuner，都有多个可选工具，差别暂不清楚。感觉，算力够的话，可以逐个 apply 选最终的 best。\n",
    "3. 官网文档：this tutorial will not run on Windows or recent versions of macOS。实测，Mac M1 确实跑不了。原因未查。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. 准备模型 & 测试数据\n",
    "\n",
    "1. 使用 Pytorch 的 AlexNet。最简单的 CNN 模型。\n",
    "2. 更大的模型也可以，但 ansor 的搜索时间更久，也不方便对比优化前后的差异，不适合学习使用。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tvm versin: 0.13.dev0\n"
     ]
    }
   ],
   "source": [
    "import tvm\n",
    "from tvm import relay\n",
    "from tvm.contrib import graph_executor\n",
    "# from tvm.ir.module import IRModule\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "\n",
    "# required by autotvm\n",
    "from tvm import autotvm\n",
    "# kernel tunners of autotvm\n",
    "from tvm.autotvm.tuner import XGBTuner, GATuner, RandomTuner, GridSearchTuner\n",
    "# graph tuners of autotvm\n",
    "from tvm.autotvm.graph_tuner import DPTuner, PBQPTuner\n",
    "\n",
    "from PIL import Image\n",
    "\n",
    "print('tvm versin: %s' % tvm.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AlexNet(\n",
       "  (features): Sequential(\n",
       "    (0): Conv2d(3, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2))\n",
       "    (1): ReLU(inplace=True)\n",
       "    (2): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (3): Conv2d(64, 192, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
       "    (4): ReLU(inplace=True)\n",
       "    (5): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "    (6): Conv2d(192, 384, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (7): ReLU(inplace=True)\n",
       "    (8): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (9): ReLU(inplace=True)\n",
       "    (10): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "    (11): ReLU(inplace=True)\n",
       "    (12): MaxPool2d(kernel_size=3, stride=2, padding=0, dilation=1, ceil_mode=False)\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=(6, 6))\n",
       "  (classifier): Sequential(\n",
       "    (0): Dropout(p=0.5, inplace=False)\n",
       "    (1): Linear(in_features=9216, out_features=4096, bias=True)\n",
       "    (2): ReLU(inplace=True)\n",
       "    (3): Dropout(p=0.5, inplace=False)\n",
       "    (4): Linear(in_features=4096, out_features=4096, bias=True)\n",
       "    (5): ReLU(inplace=True)\n",
       "    (6): Linear(in_features=4096, out_features=1000, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = torchvision.models.alexnet(weights='IMAGENET1K_V1')\n",
    "# model = torchvision.models.resnet18(weights='IMAGENET1K_V1')\n",
    "model = model.eval()\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_path = 'test-data/cat.png'\n",
    "img = Image.open(img_path).resize((224, 224))\n",
    "\n",
    "my_preprocess = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize(256),\n",
    "        transforms.CenterCrop(224),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "    ]\n",
    ")\n",
    "img = my_preprocess(img)\n",
    "# 新增Batch维度\n",
    "img = np.expand_dims(img, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export to TorchScripted model, so that we can use it in TVM\n",
    "img_input_name = 'input0'\n",
    "img_shape = [1, 3, 224, 224]\n",
    "input_shapes = [(img_input_name, img_shape)]\n",
    "\n",
    "input_data = torch.randn(img_shape)\n",
    "scripted_model = torch.jit.trace(model, input_data).eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Torch top-1 id: 285\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    torch_img = torch.from_numpy(img)\n",
    "    output = model(torch_img)\n",
    "\n",
    "    # Get top-1 result for PyTorch\n",
    "    top1_torch = np.argmax(output.numpy())\n",
    "    print('Torch top-1 id: %d' % top1_torch)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. baseline 性能 - 只编译，不优化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.\n"
     ]
    }
   ],
   "source": [
    "mod, params = relay.frontend.from_pytorch(scripted_model, input_shapes)\n",
    "\n",
    "target = 'llvm'\n",
    "# compile the model\n",
    "with tvm.transform.PassContext(opt_level=2):\n",
    "    lib = relay.build(mod, target=target, params=params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run the model\n",
    "def run_model(lib):\n",
    "    m = graph_executor.GraphModule(lib[\"default\"](tvm.cpu(0)))\n",
    "    m.set_input(img_input_name, tvm.nd.array(img.astype('float32')))\n",
    "    m.run()\n",
    "    tvm_output = m.get_output(0)\n",
    "\n",
    "    top1_tvm_1 = np.argmax(tvm_output.numpy()[0])\n",
    "    print(\"TVM top-1 id: %s\" % top1_tvm_1)\n",
    "\n",
    "    # Evaluate\n",
    "    print(\"Evaluate inference time cost...\")\n",
    "    print(m.benchmark(tvm.cpu(0), repeat=3, min_repeat_ms=500))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TVM top-1 id: 285\n",
      "Evaluate inference time cost...\n",
      "Execution time summary:\n",
      " mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  \n",
      "  42.3022      42.2435      42.4208      42.2422       0.0839                  \n"
     ]
    }
   ],
   "source": [
    "run_model(lib)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. 使用 AutoTVM 优化模型\n",
    "\n",
    "分 4 步：\n",
    "\n",
    "1. 构造 tasks. API 是 `autotvm.task.extract_from_program`。\n",
    "2. tune kernel (operator)\n",
    "3. tune graph\n",
    "4. apply best schedule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = \"llvm\"\n",
    "tasks = autotvm.task.extract_from_program(\n",
    "    mod[\"main\"], target=target, params=params, ops=(relay.op.get(\"nn.conv2d\"),)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5 <class 'tvm.autotvm.task.task.Task'>\n",
      "ConfigSpace (len=252, range_length=252, space_map=\n",
      "   0 tile_ic: Split(policy=factors, product=3, num_outputs=2) len=2\n",
      "   1 tile_oc: Split(policy=factors, product=64, num_outputs=2) len=7\n",
      "   2 tile_ow: Split(policy=verbose, product=55, num_outputs=2) len=9\n",
      "   3 unroll_kw: OtherOption([True, False]) len=2\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# example of generated tasks\n",
    "print(len(tasks), type(tasks[0]))\n",
    "print(tasks[0].config_space)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Task  1/ 5]  Current/Best:  171.69/ 545.96 GFLOPS | Progress: (30/30) | 39.33 s Done.\n",
      "[Task  2/ 5]  Current/Best:  102.52/ 589.67 GFLOPS | Progress: (30/30) | 29.11 s Done.\n",
      "[Task  3/ 5]  Current/Best:  461.77/ 530.50 GFLOPS | Progress: (30/30) | 42.68 s Done.\n",
      "[Task  4/ 5]  Current/Best:   90.55/ 365.76 GFLOPS | Progress: (30/30) | 23.37 s Done.\n",
      "[Task  5/ 5]  Current/Best:  514.72/ 514.72 GFLOPS | Progress: (30/30) | 17.84 s Done.\n"
     ]
    }
   ],
   "source": [
    "# tune kernels\n",
    "log_file = \"tune-autotvm-alexnet.json\"\n",
    "\n",
    "measure_option = autotvm.measure_option(\n",
    "    builder=autotvm.LocalBuilder(),\n",
    "    runner=autotvm.LocalRunner(\n",
    "        number=1, repeat=1, min_repeat_ms=0, enable_cpu_cache_flush=True\n",
    "    ),\n",
    ")\n",
    "\n",
    "for i, task in enumerate(tasks):\n",
    "    prefix = \"[Task %2d/%2d] \" % (i + 1, len(tasks))\n",
    "    tuner_obj = XGBTuner(task, loss_type=\"rank\")\n",
    "\n",
    "    n_trial = min(30, len(task.config_space))\n",
    "    tuner_obj.tune(\n",
    "        n_trial=n_trial,\n",
    "        early_stopping=None,\n",
    "        measure_option=measure_option,\n",
    "        callbacks=[\n",
    "            autotvm.callback.progress_bar(n_trial, prefix=prefix),\n",
    "            autotvm.callback.log_to_file(log_file),\n",
    "        ],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-06-15 17:04:28,819 INFO Start to benchmark layout transformation...\n",
      "2023-06-15 17:13:51,695 INFO Benchmarking layout transformation successful.\n",
      "2023-06-15 17:14:05,594 INFO Start to run PBQP algorithm...\n",
      "2023-06-15 17:14:05,597 INFO Finished PBQPExecutor run. Got optimal solution.\n",
      "2023-06-15 17:14:05,598 INFO Writing optimal schedules to tune-autotvm-alexnet-graph_opt.log successfully.\n"
     ]
    }
   ],
   "source": [
    "# cost 8min if use_DP = True\n",
    "# cost 10min if use_DP = False\n",
    "use_DP = False\n",
    "input_shape_dict = {img_input_name: img_shape}\n",
    "\n",
    "graph_opt_sch_file = \"tune-autotvm-alexnet-graph_opt.log\"\n",
    "\n",
    "target_op = [\n",
    "    relay.op.get(\"nn.conv2d\"),\n",
    "]\n",
    "Tuner = DPTuner if use_DP else PBQPTuner\n",
    "tuner = Tuner(mod['main'], input_shape_dict, log_file, target_op, target)\n",
    "tuner.benchmark_layout_transform(min_exec_num=2)\n",
    "tuner.run()\n",
    "tuner.write_opt_sch2record_file(graph_opt_sch_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# compile with history best schedule found by autotvm\n",
    "with autotvm.apply_history_best(log_file):\n",
    "    with tvm.transform.PassContext(opt_level=3):\n",
    "        lib2 = relay.build(mod, target=target, params=params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TVM top-1 id: 285\n",
      "Evaluate inference time cost...\n",
      "Execution time summary:\n",
      " mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  \n",
      "   8.7378       8.7448       8.7527       8.7158       0.0159                  \n"
     ]
    }
   ],
   "source": [
    "# run the tuned best model\n",
    "run_model(lib2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TVM top-1 id: 285\n",
      "Evaluate inference time cost...\n",
      "Execution time summary:\n",
      " mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  \n",
      "  42.3053      42.3333      42.3391      42.2435       0.0438                  \n"
     ]
    }
   ],
   "source": [
    "# to compare the performance, re-run model before auto-scheduler\n",
    "run_model(lib)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv-py3",
   "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.7.14"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}