2.Build_end2end_module.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gMu9UtC71BQj"
      },
      "source": [
        "本例中使用了还在开发中的relax，如果环境中已经安装配置了tvm，可能会发生冲突，建议到google colab中执行\n",
        "<a href=\"https://colab.research.google.com/drive/1vOt77ugt1AhmhufUyw-bW8emIwh0fXqN#scrollTo=PW18qc_faMIC\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mn6n9P2IlNQP"
      },
      "source": [
        "# 端到端模型执行\n",
        "\n",
        "- 关注元张量函数，讨论如何构建端到端模型\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H_f4_c9gl_pp"
      },
      "source": [
        "## 环境，数据集及模型参数准备"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3UGJvwypkwKC",
        "outputId": "bc505b2a-86be-4159-bc19-68c86c61bd51"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Looking in links: https://mlc.ai/wheels\n",
            "Requirement already satisfied: mlc-ai-nightly in /usr/local/lib/python3.7/dist-packages (0.9.dev2226+gf68b7661e)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (4.4.2)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (1.21.6)\n",
            "Requirement already satisfied: tornado in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (5.1.1)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (1.7.3)\n",
            "Requirement already satisfied: cloudpickle in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (1.5.0)\n",
            "Requirement already satisfied: synr==0.6.0 in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (0.6.0)\n",
            "Requirement already satisfied: attrs in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (22.1.0)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.7/dist-packages (from mlc-ai-nightly) (5.4.8)\n"
          ]
        }
      ],
      "source": [
        "!python3 -m  pip install mlc-ai-nightly -f https://mlc.ai/wheels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "xZMaAuz7lEjQ"
      },
      "outputs": [],
      "source": [
        "# This is needed for deferring annotation parsing in TVMScript\n",
        "from __future__ import annotations\n",
        "import IPython\n",
        "import numpy as np\n",
        "import tvm\n",
        "from tvm import relax\n",
        "from tvm.ir.module import IRModule\n",
        "from tvm.script import relax as R\n",
        "from tvm.script import tir as T"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {
        "id": "T8yRGyvJlWml"
      },
      "outputs": [],
      "source": [
        "# 加载数据集\n",
        "# 使用fashion MNIST数据集上的一个模型 \n",
        "# 以下代码从 torchvision 下载并准备数据并转换成NumPy数组。\n",
        "import torch\n",
        "import torchvision\n",
        "\n",
        "test_data = torchvision.datasets.FashionMNIST(\n",
        "    root=\"data\",\n",
        "    train=False,\n",
        "    download=True,\n",
        "    transform=torchvision.transforms.ToTensor()\n",
        ")\n",
        "test_loader = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=True)\n",
        "class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',\n",
        "               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']\n",
        "\n",
        "img, label = next(iter(test_loader))\n",
        "img = img.reshape(1, 28, 28).numpy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 287
        },
        "id": "i6oIe5-9ltfk",
        "outputId": "e3262d64-eec4-4bba-9b5d-5784cb59b193"
      },
      "outputs": [
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAS4AAAD8CAYAAADJwUnTAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAZPUlEQVR4nO3dfZAd1Znf8e9v7mgkIQmQEBYyEqDFsr2K1pEdFdiFE7OFXwRVMXYloZArG5yQyOWyEjt2UkucFKbYShW7WXCcKkIyXquMt7zGxPauFVtBaxNviLdsLOElWC/BqGQJpNULkjDoJUgz9z75417BnZc+3TP33rnd4vep6tLtPv1ybs/omT6nnz6tiMDMrEoG+l0BM7OpcuAys8px4DKzynHgMrPKceAys8px4DKzynHgMrOekbRJ0lFJOzLKJek/Sdoj6RlJ7yqyXwcuM+ulrwLrEuU3Aytb0wbgoSI7deAys56JiCeAE4lVbgW+Fk0/BS6VtDRvv4PdqmARQ5odc5g3k4c0e0N5ldOci7PqZB8f+u15cfxEvdC6Tz1zdifwatui4YgYnsLhrgReaJs/0Fp2KLVRR4FL0jrgS0AN+KOIuC+1/hzmcb1u6uSQZpbwZDze8T6On6jzs61XFVq3tvS5VyNibccHnaJpBy5JNeBB4AM0o+Q2SZsjYle3KmdmMy+ABo2ZOtxBYHnb/LLWsqRO+riuA/ZExN6IOAc8QrO9amYVFgQjUS80dcFm4B+17i6+G3g5IpLNROisqThZ2/T68StJ2kDzbgFzuKiDw5nZTOnWFZekbwA3AoslHQC+AMwCiIj/AmwBbgH2AGeAf1xkvz3vnG911A0DXKxFHkPHrOSCoN6l4a4iYn1OeQCfmup+Owlc02qbmln5NSj3NUYngWsbsFLSCpoB63bgY12plZn1TQD1CzVwRcSopI3AVprpEJsiYmfXamZmfXMhX3EREVtodq6Z2QUigJGSD+k+o5nzZlZ+QVy4TUUzu0AF1Msdtxy4zGysZuZ8uTlwmdk4ok5Hz2n3nAOXmY3R7Jx34DKzCmnmcTlwmVnFNHzFZWZV4isuM6ucQNRLPqq7A5eZTeCmoplVSiDORa3f1Uhy4DKzMZoJqG4qmlnFuHPezColQtTDV1xmVjENX3GZWZU0O+fLHRrKXTszm3HunDezSqo7j8vMqsSZ82ZWSQ3fVTSzKmk+ZO3AZWYVEogRP/JjZlUSgRNQzaxq5ARUM6uWwFdcZlZB7pw3s0oJ5IEEzaxamq8nK3doKHftzKwP/EJYM6uY4ALPnJe0DzgJ1IHRiFjbjUqZWX+V/YqrG2H1tyNijYOW2YUhQjRioNBUhKR1kp6VtEfSXZOUXyXpR5L+StIzkm7J26ebimY2RrNzvjuP/EiqAQ8CHwAOANskbY6IXW2r/Tvg0Yh4SNIqYAtwTWq/nV5xBfDnkp6StCGj4hskbZe0fYSzHR7OzHqvOeZ8kamA64A9EbE3Is4BjwC3jlsngItbny8B/jpvp51ecb03Ig5KehPwA0n/NyKeGFOjiGFgGOBiLYoOj2dmPdbsnC/cx7VY0va2+eHW//nzrgReaJs/AFw/bh/30LwA+ufAPOD9eQftKHBFxMHWv0cl/SnN6PpEeiszK7spZM4f60L/9nrgqxFxv6T3AH8saXVENLI2mHZTUdI8SQvOfwY+COyY7v7MrBzOZ84XmQo4CCxvm1/WWtbuTuBRgIj4CTAHWJzaaSd9XEuAH0v6P8DPgO9HxGMd7M/MSqLBQKGpgG3ASkkrJA0BtwObx63zPHATgKTfpBm4XkztdNpNxYjYC/zN6W5vZuUUASON7iSgRsSopI3AVqAGbIqInZLuBbZHxGbgc8CXJf1Lml1sH4+IZH+40yHMbIxmU7F7mfMRsYVmikP7srvbPu8CbpjKPh24zGyCsmfOO3CZ2RhTTIfoCwcuMxunu03FXnDgMrMJPOa8VZtyfoHTN3+SGu9dkywf+PHT0953rh5+r6pr3lX068nMrEI8dLOZVZKbimZWKb6raGaV5LuKZlYpEWLUgcvMqsZNRTOrFPdxWelp9uxkeZztbLjtfd98R2bZ2644nNz21O/9rWT5rB8+Na06AW/oPK0iHLjMrFKcx2VmleQ8LjOrlAgY7dJAgr3iwGVmE7ipaGaV4j4uM6ukcOAys6px57z1XmpsqZx8pU7ztPLywC793rzMspdOX53c9sDt9fSx/+67k+UrP/3TzLKBOXOS28boaEfleQavWJJZVj92vKfHzhPhPi4zqxxR911FM6sa93GZWaX4WUUzq54o/6OcDlxmNoHvKppZpYQ7582sitxUtN6/wy+xfe2yRclNR1anc6me/0A636kxlCzmk7dszSx78On3Jbe95Gdzk+Vvfjyd75TKAmu8+mpy216rL7s8s6w2O31SR/e/0O3qTFD2u4q514OSNkk6KmlH27JFkn4g6bnWvwt7W00zmykRzcBVZOqXIg3ZrwLrxi27C3g8IlYCj7fmzewC0QgVmvolN3BFxBPAiXGLbwUebn1+GPhIl+tlZn0UUWzql+n2cS2JiEOtz4eBzAevJG0ANgDM4aJpHs7MZkogGiW/q9hx7SIiaCbbZpUPR8TaiFg7i/QDuWZWDlFw6pfpBq4jkpYCtP492r0qmVlfdblzXtI6Sc9K2iNp0v5wSbdJ2iVpp6Q/ydvndAPXZuCO1uc7gO9Ocz9mVkZduuSSVAMeBG4GVgHrJa0at85K4N8AN0TE3wA+k7ff3D4uSd8AbgQWSzoAfAG4D3hU0p3AfuC2/K9QcR2MedXrXswXP/mezLK5x9PH/usPpcd2mrs3/Vd19Jr/lyxP5Wp96G27k9suWn06Wf6dS/52svyaw9k5bPXj4+83zazYviOzrLejbRXTxVSH64A9EbEXQNIjNG/u7Wpb558BD0bES81jR24LLjdwRcT6jKKb8rY1s+oJoNEoHLgWS9reNj8cEcNt81cC7RmzB4Drx+3jrQCS/hKoAfdExGOpgzpz3szGCqD4FdexiFjb4REHgZU0W3bLgCck/VZE/Dprg3Lf8zSzvuhiHtdBYHnb/LLWsnYHgM0RMRIRvwJ+STOQZXLgMrOJupcPsQ1YKWmFpCHgdpo399r9Gc2rLSQtptl03JvaqZuKZjZO955DjIhRSRuBrTT7rzZFxE5J9wLbI2Jzq+yDknbRfDb+X0dE8gl6By4zm6iLN8IjYguwZdyyu9s+B/DZ1lRIuQJXr4d/6UQHx65dfHGy/OT7fzNZfmZxukX/6qLs83bqquSmDB2eld730vQrwi666FyyfMHc7OFjnj+dHlRkycJXkuW7Pvmfk+WP3ZH9pMbjr6zKLAP4799Lv/rs6rt/kiyvvWVFsvyFjy7NLDu9rJHcNvXata4IiOJ3FfuiXIHLzErCgcvMqsYjoJpZ5ThwmVmlTC0BtS8cuMxsAr8sw8yqx3cVzaxq5CuuKejj9akG06dCq9+aWXZm2fzktoffU0uWD55K/3U7c206V6r2cnbdh64+ldz27W86kiy/dv6xZPmaec8nyy+rZR//lUb61WfPnEknof3ukTXJ8hWzX8wsWzk3/b0/+w/SQ8xd9bH0q9EeS/xMAN5K9ivG5tfOJrd95P4bMsvOPtCFHK9+D29aQLkCl5mVgNw5b2YV5CsuM6uc9FNHfefAZWZjOY/LzKrIdxXNrHpKHrg8AqqZVU6prrhql1+eLB9525WZZS9fOze57Yl3pP+EXLo73aZ/+S3ZZaOXjyS3ZSA9plXMSb+Q6uK56TyuVxdkj6k1Zyhdt5Mj6Vyq/WeyX/EFsO3Y1cnyY6fmZZYNDaa/d20g/TNbvfhQsvzkaPZ3+/W59O/LpUPp167tGnhzsvz0aPqt7cfPXpRZdslQ9hhmABevfCmz7HDO71JRbiqaWbUEfuTHzCrIV1xmVjVuKppZ9ThwmVnlOHCZWZUo3FQ0syryXcXXae4cBt769szySx9Kj5H0k92J6p5L50oNnkqPiTV6Ud4PKvtP0OCL6XcTji5O51KNnkvX7WROzpFS+U7pTXMN5eSgrVp4OFm+6E2nM8uump0e02reQHpcqjlKn9d64hVbedsOKf298ywYSOeBzUrs/7lzVyS3nVvLrvvBWemcv6LKfsWVmzkvaZOko5J2tC27R9JBSU+3plt6W00zm1FRcOqTIo/8fBVYN8nyL0bEmta0ZZJyM6uieL2fK2/ql9zAFRFPACdmoC5mVhYXwBVXlo2Snmk1JRdmrSRpg6TtkrafG83u7zCz8lCj2NQv0w1cDwHXAmuAQ8D9WStGxHBErI2ItUOD2Q/cmpkVNa3AFRFHIqIeEQ3gy8B13a2WmfXVhdhUlLS0bfajwI6sdc2sYirQOZ+bxyXpG8CNwGJJB4AvADdKWkMz5u4DPlHoaCMj6ODRzOJnH87O8QKYvSy77NzV6Zyf2Vek+9cG3p4ex+hfXJv9vrrvH1md3PY9l/0qWb7tpfSYVnljQ/3WgoOZZW+Znc6Nu2ZW+r2JeWbldHTsG8kez+svT2W/qxLghTOZXacAvJIzltgvD70ps2zkTDr3Lu9qYuBk+r/Ogl+lrwkiUTyQM7zbkidfySw7uz89DlhhJc/jyg1cEbF+ksVf6UFdzKwsqh64zOyNRfT3jmERHnPezMbqch+XpHWSnpW0R9JdifX+nqSQtDZvnw5cZjZRl+4qSqoBDwI3A6uA9ZJWTbLeAuDTwJNFqufAZWYTdS8d4jpgT0TsjYhzwCPArZOs93vA7wPpN4W0OHCZ2QRTaCouPv9kTGvaMG5XVwIvtM0faC17/VjSu4DlEfH9ovWb2c75Wg0WXpJZvOTbzyY3rx9PPDKp9LA0A7PTt4k1L/t1UQCPzc1OeajV0vH/yXnpdIkYTG9/fCCdFvC/RpZkl0X6z6JOp1Mt4lTOY1oj6TSS+smTiZ3n/cnOfg0XgHJ+pr/RwSO2A/NznvIYSA9FlEdDiXSMVBkQF2WngajRpduBxXdzLCJy+6SySBoAHgA+PpXtfFfRzMaKrt5VPAgsb5tf1lp23gJgNfAXal58XAFslvThiNietVMHLjObqHt5XNuAlZJW0AxYtwMfe+0wES8Di8/PS/oL4F+lgha4j8vMJtGtdIiIGAU2AluB3cCjEbFT0r2SPjzd+vmKy8wm6mLmfGug0S3jlt2dse6NRfbpwGVmY/V55IciHLjMbAxR/pdlOHCZ2QQOXO2kZI5KY8Wbk5s3VmcP/zL4cnpYm7xcKZ3MSdhN5TudzXkl1Jn0vgdy8sCo5eQMDSbKB9L7bixckCyvvzl7WBqAGErvvz4ru7w+J/29Iudr571BrJOcpoFzneUDRC2dVxgD2eWNofS2g6ezv3hjf2f5Za9x4DKzynHgMrNK6fPopkU4cJnZRA5cZlY1ZR9I0IHLzCZwU9HMqsUJqGZWSQ5cbSRiKPuQGk03rGcd/HV24eyh9LEb6X3H7PSpaMy/NHvbnDysmNVZbk0jJ1eqkcgZytu2djZ9XgZyfiZ5aiPZ29deTSdiqZ73jrB0vlMvRd6hc/K4UkYb6d+XoZey8wJV77xzypnzZlZJXRuQsEccuMxsLPdxmVkVualoZtXjwGVmVeMrLjOrHgcuM6uU7r7lpydyA5ek5cDXgCU04/BwRHxJ0iLgm8A1wD7gtohIvwiv0UCns3NQzi3LzpUCYGH2++Ty5I2vNHA2/X7AlNxbxyPpfKW8PLCBnHyngUQ+UyTyqGZEIhcrrzkykHPeGjn5camxwlLjYUE6Nw7y695JOkF9Tvr3YeTS7P8Heb9LRVQhj6vItxwFPhcRq4B3A5+StAq4C3g8IlYCj7fmzexCEFFs6pPcwBURhyLi563PJ2m+YuhK4Fbg4dZqDwMf6VUlzWxmdev1ZL0ypT4uSdcA7wSeBJZExKFW0WGaTUkzq7oLKQFV0nzg28BnIuKV1uuyAYiIkCaPv5I2ABsA5gxe3FltzWxGlL1zvlBPnqRZNIPW1yPiO63FRyQtbZUvBY5Otm1EDEfE2ohYO1Sb2406m1mPqVFs6pfcwKXmpdVXgN0R8UBb0WbgjtbnO4Dvdr96ZjbjgtJ3zhdpKt4A/A7wC0lPt5Z9HrgPeFTSncB+4La8HcXZc9Sf25tZXnsuvX3tsuxXZWlu+mqufsXC9M5zNBLD3uQNHTOa8xquWk6qRn12TrpEB6/SyutgbQym0wJy65747nnD1sT8ztIMBxNpJCPz0vuee+h0snz04tnJ8rz9zzqdnX4z51j6dXeDu/Znlg2cSb+mr6iyp0Pk/mZExI9ppnZM5qbuVsfMSqHqgcvM3liqkIDqwGVmY0V4IEEzq6Byxy0HLjObyE1FM6uWANxUNLPKKXfcqlbgqh8/Mf2NDxzs6NipbKa8l4919nKyiv2Qxun0u/dK3gBJef9ve/0zT0kN9hORHgqoqG42FSWtA75E87T8UUTcN678s8A/pTkSzYvAP4mI7GQ1Cj7yY2ZvLGpEoSl3P1INeBC4GVgFrG8Ni9Xur4C1EfEO4FvAH+Tt14HLzMaKKUz5rgP2RMTeiDgHPEJzSKzXDxfxo4g405r9KbAsb6dVboWYWQ80E1ALtxUXS9reNj8cEcNt81cCL7TNHwCuT+zvTuB/5B3UgcvMJir++OuxiFjbjUNK+ofAWuB9ees6cJnZBFO44spzEFjeNr+stWzs8aT3A/8WeF9E5D4p7j4uMxuru31c24CVklZIGgJupzkk1mskvRP4r8CHI2LScf3G8xWXmY3TvWcVI2JU0kZgK810iE0RsVPSvcD2iNgM/AdgPvDfWiMrPx8RH07t14HLzCbq4iCBEbEF2DJu2d1tn98/1X06cJnZWBfCC2HN7A2oj8MyF+HAZWYTlTtuOXCZ2URqlLut6MBlZmMFU0lA7QsHLjMbQ0Q3E1B7woHLzCZy4DKzynHgMrNKcR+XmVWR7yqaWcWEm4pmVjGBA5eZVVC5W4oOXGY2kfO4zKx6Sh64ckdAlbRc0o8k7ZK0U9KnW8vvkXRQ0tOt6ZbeV9fMei4C6o1iU58UueIaBT4XET+XtAB4StIPWmVfjIg/7F31zKwvSn7FlRu4IuIQcKj1+aSk3TRfOWRmF6qSB64pvSxD0jXAO4EnW4s2SnpG0iZJCzO22SBpu6TtI+S+vMPM+i2ARhSb+qRw4JI0H/g28JmIeAV4CLgWWEPziuz+ybaLiOGIWBsRa2cxuwtVNrPeCohGsalPCt1VlDSLZtD6ekR8ByAijrSVfxn4Xk9qaGYzK+hrx3sRRe4qCvgKsDsiHmhbvrRttY8CO7pfPTPri4hiU58UueK6Afgd4BeSnm4t+zywXtIamvF5H/CJntTQzGZeyTvni9xV/DGgSYq2TLLMzCrPD1mbWdUE4GFtzKxyfMVlZtUSpb+r6MBlZmMFRB9ztIpw4DKzifqYFV+EA5eZTeQ+LjOrlAjfVTSzCvIVl5lVSxD1er8rkeTAZWZjnR/WpsQcuMxsopKnQ0xpIEEzu/AFEI0oNBUhaZ2kZyXtkXTXJOWzJX2zVf5ka8DSJAcuMxsrujeQoKQa8CBwM7CK5qgyq8atdifwUkS8Bfgi8Pt5+3XgMrMJol4vNBVwHbAnIvZGxDngEeDWcevcCjzc+vwt4KbWOICZZrSP6yQvHfthfGt/26LFwLGZrMMUlLVuZa0XuG7T1c26Xd3pDk7y0tYfxrcWF1x9jqTtbfPDETHcNn8l8ELb/AHg+nH7eG2diBiV9DJwGYlzMqOBKyIub5+XtD0i1s5kHYoqa93KWi9w3aarbHWLiHX9rkMeNxXNrJcOAsvb5pe1lk26jqRB4BLgeGqnDlxm1kvbgJWSVkgaAm4HNo9bZzNwR+vz3wf+Z0Q6db/feVzD+av0TVnrVtZ6ges2XWWuW0dafVYbga1ADdgUETsl3Qtsj4jNNF/G88eS9gAnaAa3JOUENjOz0nFT0cwqx4HLzCqnL4Er7xGAfpK0T9IvJD09Lj+lH3XZJOmopB1tyxZJ+oGk51r/LixR3e6RdLB17p6WdEuf6rZc0o8k7ZK0U9KnW8v7eu4S9SrFeauSGe/jaj0C8EvgAzST0bYB6yNi14xWJIOkfcDaiOh7sqKkvwOcAr4WEatby/4AOBER97WC/sKI+N2S1O0e4FRE/OFM12dc3ZYCSyPi55IWAE8BHwE+Th/PXaJet1GC81Yl/bjiKvIIgAER8QTNuyzt2h+PeJjmL/6My6hbKUTEoYj4eevzSWA3zezsvp67RL1sivoRuCZ7BKBMP7wA/lzSU5I29Lsyk1gSEYdanw8DS/pZmUlslPRMqynZl2Zsu9ZIA+8EnqRE525cvaBk563s3Dk/0Xsj4l00n2b/VKtJVEqtJL0y5bM8BFwLrAEOAff3szKS5gPfBj4TEa+0l/Xz3E1Sr1KdtyroR+Aq8ghA30TEwda/R4E/pdm0LZMjrb6S830mR/tcn9dExJGIqEfzpXxfpo/nTtIsmsHh6xHxndbivp+7yepVpvNWFf0IXEUeAegLSfNanaZImgd8ENiR3mrGtT8ecQfw3T7WZYzzQaHlo/Tp3LWGRPkKsDsiHmgr6uu5y6pXWc5blfQlc751u/c/8vojAP9+xisxCUm/QfMqC5qPQ/1JP+sm6RvAjTSHPTkCfAH4M+BR4CpgP3BbRMx4J3lG3W6k2dwJYB/wibY+pZms23uB/w38Ajg/2t3nafYn9e3cJeq1nhKctyrxIz9mVjnunDezynHgMrPKceAys8px4DKzynHgMrPKceAys8px4DKzyvn//PmBlqIwC8AAAAAASUVORK5CYII=",
            "text/plain": [
              "<Figure size 432x288 with 2 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Class: Sneaker\n"
          ]
        }
      ],
      "source": [
        "# 绘制我们希望能预测的图像\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "plt.figure()\n",
        "plt.imshow(img[0])\n",
        "plt.colorbar()\n",
        "plt.grid(False)\n",
        "plt.show()\n",
        "print(\"Class:\", class_names[label[0]])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PgkkuxRVl3Fk",
        "outputId": "757efab2-7b1b-4506-cffe-7103d625147b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--2022-08-23 10:13:58--  https://github.com/mlc-ai/web-data/raw/main/models/fasionmnist_mlp_params.pkl\n",
            "Resolving github.com (github.com)... 140.82.114.4\n",
            "Connecting to github.com (github.com)|140.82.114.4|:443... connected.\n",
            "HTTP request sent, awaiting response... 302 Found\n",
            "Location: https://raw.githubusercontent.com/mlc-ai/web-data/main/models/fasionmnist_mlp_params.pkl [following]\n",
            "--2022-08-23 10:13:58--  https://raw.githubusercontent.com/mlc-ai/web-data/main/models/fasionmnist_mlp_params.pkl\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.109.133, 185.199.108.133, 185.199.111.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.109.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 407396 (398K) [application/octet-stream]\n",
            "Saving to: ‘fasionmnist_mlp_params.pkl.1’\n",
            "\n",
            "fasionmnist_mlp_par 100%[===================>] 397.85K  --.-KB/s    in 0.04s   \n",
            "\n",
            "2022-08-23 10:13:59 (10.8 MB/s) - ‘fasionmnist_mlp_params.pkl.1’ saved [407396/407396]\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# 下载模型参数\n",
        "# Hide outputs\n",
        "!wget https://github.com/mlc-ai/web-data/raw/main/models/fasionmnist_mlp_params.pkl"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GPS1VfL4mHpc"
      },
      "source": [
        "## 端到端模型整合\n",
        "\n",
        "例子：一个两层神经网络，由两个全连接层和一个有 relu 激活层组成。为了简化问题，我们删除了最终的 softmax 层。输出分数是未标准化的，但最大值仍然对应于最可能的类别。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "U_4xHrQDmK35",
        "outputId": "a3a288c3-a3cf-40ee-93f9-95d6342e1a8d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[[-18.002035  -26.690979  -14.855994  -17.64491   -21.919956   -3.1951513\n",
            "  -18.476469   17.237576   -0.7316637 -11.115718 ]]\n",
            "[7]\n",
            "NumPy-MLP Prediction: Sneaker\n"
          ]
        }
      ],
      "source": [
        "# 1.numpy实现\n",
        "def numpy_mlp(data, w0, b0, w1, b1):\n",
        "    lv0 = data @ w0.T + b0\n",
        "    lv1 = np.maximum(lv0, 0)\n",
        "    lv2 = lv1 @ w1.T + b1\n",
        "    return lv2\n",
        "\n",
        "import pickle as pkl\n",
        "\n",
        "mlp_params = pkl.load(open(\"fasionmnist_mlp_params.pkl\", \"rb\"))\n",
        "res = numpy_mlp(img.reshape(1, 784),\n",
        "                mlp_params[\"w0\"],\n",
        "                mlp_params[\"b0\"],\n",
        "                mlp_params[\"w1\"],\n",
        "                mlp_params[\"b1\"])\n",
        "print(res)\n",
        "pred_kind = res.argmax(axis=1)\n",
        "print(pred_kind)\n",
        "print(\"NumPy-MLP Prediction:\", class_names[pred_kind[0]])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FDVi_I4KJl55",
        "outputId": "58f83b7b-0249-4b01-e775-dcfc57a3f02d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Low-level Numpy MLP Prediction: Sneaker\n"
          ]
        }
      ],
      "source": [
        "# 2.low-level numpy 实现\n",
        "def lnumpy_linear0(X: np.ndarray, W: np.ndarray, B: np.ndarray, Z: np.ndarray):\n",
        "    Y = np.empty((1, 128), dtype=\"float32\")\n",
        "    for i in range(1):\n",
        "        for j in range(128):\n",
        "            for k in range(784):\n",
        "                if k == 0:\n",
        "                    Y[i, j] = 0\n",
        "                Y[i, j] = Y[i, j] + X[i, k] * W[j, k]\n",
        "\n",
        "    for i in range(1):\n",
        "        for j in range(128):\n",
        "            Z[i, j] = Y[i, j] + B[j]\n",
        "\n",
        "\n",
        "def lnumpy_relu0(X: np.ndarray, Y: np.ndarray):\n",
        "     for i in range(1):\n",
        "        for j in range(128):\n",
        "            Y[i, j] = np.maximum(X[i, j], 0)\n",
        "\n",
        "def lnumpy_linear1(X: np.ndarray, W: np.ndarray, B: np.ndarray, Z: np.ndarray):\n",
        "    Y = np.empty((1, 10), dtype=\"float32\")\n",
        "    for i in range(1):\n",
        "        for j in range(10):\n",
        "            for k in range(128):\n",
        "                if k == 0:\n",
        "                    Y[i, j] = 0\n",
        "                Y[i, j] = Y[i, j] + X[i, k] * W[j, k]\n",
        "\n",
        "    for i in range(1):\n",
        "        for j in range(10):\n",
        "            Z[i, j] = Y[i, j] + B[j]\n",
        "\n",
        "\n",
        "def lnumpy_mlp(data, w0, b0, w1, b1):\n",
        "    lv0 = np.empty((1, 128), dtype=\"float32\")\n",
        "    lnumpy_linear0(data, w0, b0, lv0)\n",
        "\n",
        "    lv1 = np.empty((1, 128), dtype=\"float32\")\n",
        "    lnumpy_relu0(lv0, lv1)\n",
        "\n",
        "    out = np.empty((1, 10), dtype=\"float32\")\n",
        "    lnumpy_linear1(lv1, w1, b1, out)\n",
        "    return out\n",
        "\n",
        "result =lnumpy_mlp(\n",
        "    img.reshape(1, 784),\n",
        "    mlp_params[\"w0\"],\n",
        "    mlp_params[\"b0\"],\n",
        "    mlp_params[\"w1\"],\n",
        "    mlp_params[\"b1\"])\n",
        "\n",
        "pred_kind = result.argmax(axis=1)\n",
        "print(\"Low-level Numpy MLP Prediction:\", class_names[pred_kind[0]])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o6hl9dZuJ4Pc"
      },
      "source": [
        "### 在TVMScript中构建端到端IRModule\n",
        "- 参考Numpy代码，用MLC抽象实现端到端模型运行\n",
        "\n",
        "- R.function 为relax函数。relax函数是一种表示上层神经网络执行的全新抽象"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "id": "ugBQOYSFKB2B"
      },
      "outputs": [],
      "source": [
        "@tvm.script.ir_module\n",
        "class MyModule:\n",
        "    @T.prim_func\n",
        "    def relu0(X: T.Buffer[(1, 128), \"float32\"],\n",
        "              Y: T.Buffer[(1, 128), \"float32\"]):\n",
        "        # function attr dict\n",
        "        T.func_attr({\"global_symbol\": \"relu0\", \"tir.noalias\": True})\n",
        "        for i, j in T.grid(1, 128):\n",
        "            with T.block(\"Y\"):\n",
        "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
        "                Y[vi, vj] = T.max(X[vi, vj], T.float32(0))\n",
        "\n",
        "    @T.prim_func\n",
        "    def linear0(X: T.Buffer[(1, 784), \"float32\"],\n",
        "                W: T.Buffer[(128, 784), \"float32\"],\n",
        "                B: T.Buffer[(128,), \"float32\"],\n",
        "                Z: T.Buffer[(1, 128), \"float32\"]):\n",
        "        T.func_attr({\"global_symbol\": \"linear0\", \"tir.noalias\": True})\n",
        "        Y = T.alloc_buffer((1, 128), \"float32\")\n",
        "        for i, j, k in T.grid(1, 128, 784):\n",
        "            with T.block(\"Y\"):\n",
        "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
        "                with T.init():\n",
        "                    Y[vi, vj] = T.float32(0)\n",
        "                Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]\n",
        "\n",
        "        for i, j in T.grid(1, 128):\n",
        "            with T.block(\"Z\"):\n",
        "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
        "                Z[vi, vj] =  Y[vi, vj] + B[vj]\n",
        "\n",
        "    @T.prim_func\n",
        "    def linear1(X: T.Buffer[(1, 128), \"float32\"],\n",
        "                W: T.Buffer[(10, 128), \"float32\"],\n",
        "                B: T.Buffer[(10,), \"float32\"],\n",
        "                Z: T.Buffer[(1, 10), \"float32\"]):\n",
        "        T.func_attr({\"global_symbol\": \"linear1\", \"tir.noalias\": True})\n",
        "        Y = T.alloc_buffer((1, 10), \"float32\")\n",
        "        for i, j, k in T.grid(1, 10, 128):\n",
        "            with T.block(\"Y\"):\n",
        "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
        "                with T.init():\n",
        "                    Y[vi, vj] = T.float32(0)\n",
        "                Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]\n",
        "\n",
        "        for i, j in T.grid(1, 10):\n",
        "            with T.block(\"Z\"):\n",
        "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
        "                Z[vi, vj] = Y[vi, vj] + B[vj]\n",
        "\n",
        "    @R.function\n",
        "    def main(x: Tensor((1, 784), \"float32\"),\n",
        "             w0: Tensor((128, 784), \"float32\"),\n",
        "             b0: Tensor((128,), \"float32\"),\n",
        "             w1: Tensor((10, 128), \"float32\"),\n",
        "             b1: Tensor((10,), \"float32\")):\n",
        "        with R.dataflow():\n",
        "            lv0 = R.call_tir(linear0, (x, w0, b0), (1, 128), dtype=\"float32\")\n",
        "            lv1 = R.call_tir(relu0, (lv0,), (1, 128), dtype=\"float32\")\n",
        "            out = R.call_tir(linear1, (lv1, w1, b1), (1, 10), dtype=\"float32\")\n",
        "            R.output(out)\n",
        "        return out\n",
        "# print(MyModule.script())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SLesjqyDPMea"
      },
      "source": [
        "- 使用计算图来可视化高层模型执行\n",
        "  - 图中的框表示计算操作\n",
        "  - 箭头对应于中间张量的输入输出\n",
        "\n",
        "- 计算图通常具有的特性：\n",
        "  - 框的每个输入边对应操作的输入\n",
        "  - 每个出边对应于操作的输出\n",
        "  - 每个操作可以任意重新排序，直到边缘的拓扑排序\n",
        "\n",
        "- call_tir\n",
        "  - R.call_tir是引入元张量函数的过程。\n",
        "  - 将输入(n)和输出(1)传递给prim_func。 执行prim_func后，结果会填充到输出中。\n",
        "  - 目标传递风格：\n",
        "    - 输入和输出在外部显式分配并传递给底层函数，常用于底层库设计，便于高层框架处理内存分配决策。\n",
        "    - 适用于<font color=red>输出固定</font>的张量操作。\n",
        "  - 对于计算图的特性来说\n",
        "    - 我们要隐藏可能的分配或对函数的显式写入，及pure或side-effect free。\n",
        "    - 如果一个函数只从其输入中读取并通过其输出返回结果，它不会改变程序的其他部分（例如递增全局计数器），那么它是pure或side-effect free的。\n",
        "    - call_tir使我们能隐藏调用低层元函数细节，并将它们应用到计算图中。\n",
        "\n",
        "- 实际上，最底层的实现会有显式的内存分配，所以call_tir主要是为了让我们在生成实际实现之前继续做一些高层的转换。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Feu6B93P2NCg",
        "outputId": "448204c9-2a46-4d01-e9fe-a0db497a5bd0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Low-level Numpy with CallTIR Prediction: Sneaker\n"
          ]
        }
      ],
      "source": [
        "# 与之相对应的low-level numpy实现\n",
        "def lnumpy_call_tir(prim_func, inputs, shape, dtype):\n",
        "    res = np.empty(shape, dtype=dtype)\n",
        "    prim_func(*inputs, res)\n",
        "    return res\n",
        "\n",
        "def lnumpy_mlp_with_call_tir(data, w0, b0, w1, b1):\n",
        "    lv0 = lnumpy_call_tir(lnumpy_linear0, (data, w0, b0), (1, 128), dtype=\"float32\")\n",
        "    lv1 = lnumpy_call_tir(lnumpy_relu0, (lv0, ), (1, 128), dtype=\"float32\")\n",
        "    out = lnumpy_call_tir(lnumpy_linear1, (lv1, w1, b1), (1, 10), dtype=\"float32\")\n",
        "    return out\n",
        "\n",
        "result = lnumpy_mlp_with_call_tir(\n",
        "    img.reshape(1, 784),\n",
        "    mlp_params[\"w0\"],\n",
        "    mlp_params[\"b0\"],\n",
        "    mlp_params[\"w1\"],\n",
        "    mlp_params[\"b1\"])\n",
        "\n",
        "pred_kind = np.argmax(result, axis=1)\n",
        "print(\"Low-level Numpy with CallTIR Prediction:\", class_names[pred_kind[0]])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HawA4dqw3Fei"
      },
      "source": [
        "- Dataflow Block\n",
        "  - R.dataflow() 范围标注，是用来标记程序计算图区域的一种方式\n",
        "  - 在dataflow block中，所有操作都需要side-effect free；而dataflow block外，操作可能包含side-effect。\n",
        "    - 例如在R.dataflow()前用R.alloc_tensor创建一个张量，这个张量可以与main的参数张量一同使用。\n",
        "    - 本课的大部分课程只会处理计算图。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HrQ6WBmy6vab"
      },
      "source": [
        "## 构建并运行模型\n",
        "！注意：relax仍在开发中，因此某些API可能会更改\n",
        "- 使用relax.vm.build构建该函数\n",
        "  - 构建一个“可执行文件”。区别于传统操作系统中的可执行文件，不能直接在系统中运行，而是针对Relax VM设计的一种文件格式。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h-kYkk6R6yzB",
        "outputId": "a8d532a1-3f08-46fd-d860-d198750ab1c9"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "<class 'tvm.relax.vm.Executable'>\n",
            "[[-18.002035  -26.690985  -14.8559885 -17.64491   -21.919956   -3.1951506\n",
            "  -18.476469   17.237574   -0.731666  -11.115719 ]]\n"
          ]
        }
      ],
      "source": [
        "# 使用realx.vm.build构建一个可执行文件\n",
        "ex = relax.vm.build(MyModule, target=\"llvm\")\n",
        "print(type(ex))\n",
        "\n",
        "# 初始化一个虚拟机执行器，用来运行该函数\n",
        "vm = relax.VirtualMachine(ex, tvm.cpu())  # tvm.cpu()指示我们要在哪个设备上运行端到端执行。\n",
        "\n",
        "# 运行\n",
        "data_nd = tvm.nd.array(img.reshape(1, 784))\n",
        "nd_params = {k: tvm.nd.array(v) for k, v in mlp_params.items()}\n",
        "\n",
        "nd_res = vm[\"main\"](data_nd,\n",
        "                    nd_params[\"w0\"],\n",
        "                    nd_params[\"b0\"],\n",
        "                    nd_params[\"w1\"],\n",
        "                    nd_params[\"b1\"])\n",
        "print(nd_res)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OwFlgQq28Syl",
        "outputId": "eb3b1b4d-6bd7-4971-e3ac-0c954ca27b42"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "MyModule Prediction: Sneaker\n"
          ]
        }
      ],
      "source": [
        "# 将结果转成numpy数组，并调用argmax获取类标签。\n",
        "pred_kind = np.argmax(nd_res.numpy(), axis=1)\n",
        "print(\"MyModule Prediction:\", class_names[pred_kind[0]])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jNcRslC68qFH"
      },
      "source": [
        "## 在环境中集成现有运行库\n",
        "将现有的库函数集成到 MLC 过程中"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "id": "V_l0qq_D99iq"
      },
      "outputs": [],
      "source": [
        "@tvm.script.ir_module\n",
        "class MyModuleWithExternCall:\n",
        "    @R.function\n",
        "    def main(x: Tensor((1, 784), \"float32\"),\n",
        "             w0: Tensor((128, 784), \"float32\"),\n",
        "             b0: Tensor((128,), \"float32\"),\n",
        "             w1: Tensor((10, 128), \"float32\"),\n",
        "             b1: Tensor((10,), \"float32\")):\n",
        "        # block 0\n",
        "        with R.dataflow():\n",
        "            lv0 = R.call_tir(\"env.linear\", (x, w0, b0), (1, 128), dtype=\"float32\")\n",
        "            lv1 = R.call_tir(\"env.relu\", (lv0,), (1, 128), dtype=\"float32\")\n",
        "            out = R.call_tir(\"env.linear\", (lv1, w1, b1), (1, 10), dtype=\"float32\")\n",
        "            R.output(out)\n",
        "        return out"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BExOzu9HMkfY"
      },
      "source": [
        "### 注册运行时函数\n",
        "- DLPack 是一种通用的交换标准，允许不同的框架交换 Tensor/NDArray 而无需参与数据复制。\n",
        "\n",
        "- from_dlpack 是一个零拷贝转换，torch阵列和TVM NDArray共享底层内存。\n",
        "\n",
        "- from_dlpack API 由多个框架支持，是 Python 数组 API 标准的一部分。 如果您有兴趣，可以阅读更多 [这里](https://dmlc.github.io/dlpack/latest/python_spec.html)。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "id": "Fil0bHsEKDKR"
      },
      "outputs": [],
      "source": [
        "@tvm.register_func(\"env.linear\", override=True)\n",
        "def torch_linear(x: tvm.nd.NDArray,\n",
        "                 w: tvm.nd.NDArray,\n",
        "                 b: tvm.nd.NDArray,\n",
        "                 out: tvm.nd.NDArray):\n",
        "    x_torch = torch.from_dlpack(x)\n",
        "    w_torch = torch.from_dlpack(w)\n",
        "    b_torch = torch.from_dlpack(b)\n",
        "    out_torch = torch.from_dlpack(out)\n",
        "    torch.mm(x_torch, w_torch.T, out=out_torch)\n",
        "    torch.add(out_torch, b_torch, out=out_torch)\n",
        "\n",
        "@tvm.register_func(\"env.relu\", override=True)\n",
        "def lnumpy_relu(x: tvm.nd.NDArray,\n",
        "                out: tvm.nd.NDArray):\n",
        "    x_torch = torch.from_dlpack(x)\n",
        "    out_torch = torch.from_dlpack(out)\n",
        "    torch.maximum(x_torch, torch.Tensor([0.0]), out=out_torch)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DLVXSUr9OMHo"
      },
      "source": [
        "### build and run"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LeJgu3pmOJe6",
        "outputId": "03f8c657-de6e-42bf-ce90-1bf25f9278c6"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "MyModuleWithExternCall Prediction: Sneaker\n"
          ]
        }
      ],
      "source": [
        "ex = relax.vm.build(MyModuleWithExternCall, target=\"llvm\")\n",
        "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
        "\n",
        "nd_res = vm[\"main\"](data_nd,\n",
        "                    nd_params[\"w0\"],\n",
        "                    nd_params[\"b0\"],\n",
        "                    nd_params[\"w1\"],\n",
        "                    nd_params[\"b1\"])\n",
        "\n",
        "pred_kind = np.argmax(nd_res.numpy(), axis=1)\n",
        "print(\"MyModuleWithExternCall Prediction:\", class_names[pred_kind[0]])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zkWqxHjOOkWS"
      },
      "source": [
        "### 混合起来(TensorIR实现和库实现)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z-wtBHjbOsCW",
        "outputId": "059eb38b-3bda-4a67-a36f-3a12beb9178d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "MyModuleMixture Prediction: Sneaker\n"
          ]
        }
      ],
      "source": [
        "@tvm.script.ir_module\n",
        "class MyModuleMixture:\n",
        "    @T.prim_func\n",
        "    def linear0(X: T.Buffer[(1, 784), \"float32\"],\n",
        "                W: T.Buffer[(128, 784), \"float32\"],\n",
        "                B: T.Buffer[(128,), \"float32\"],\n",
        "                Z: T.Buffer[(1, 128), \"float32\"]):\n",
        "        T.func_attr({\"global_symbol\": \"linear0\", \"tir.noalias\": True})\n",
        "        Y = T.alloc_buffer((1, 128), \"float32\")\n",
        "        for i, j, k in T.grid(1, 128, 784):\n",
        "            with T.block(\"Y\"):\n",
        "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
        "                with T.init():\n",
        "                    Y[vi, vj] = T.float32(0)\n",
        "                Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]\n",
        "\n",
        "        for i, j in T.grid(1, 128):\n",
        "            with T.block(\"Z\"):\n",
        "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
        "                Z[vi, vj] =  Y[vi, vj] + B[vj]\n",
        "\n",
        "    @R.function\n",
        "    def main(x: Tensor((1, 784), \"float32\"),\n",
        "             w0: Tensor((128, 784), \"float32\"),\n",
        "             b0: Tensor((128,), \"float32\"),\n",
        "             w1: Tensor((10, 128), \"float32\"),\n",
        "             b1: Tensor((10,), \"float32\")):\n",
        "        with R.dataflow():\n",
        "            lv0 = R.call_tir(linear0, (x, w0, b0), (1, 128), dtype=\"float32\")\n",
        "            lv1 = R.call_tir(\"env.relu\", (lv0,), (1, 128), dtype=\"float32\")\n",
        "            out = R.call_tir(\"env.linear\", (lv1, w1, b1), (1, 10), dtype=\"float32\")\n",
        "            R.output(out)\n",
        "        return out\n",
        "\n",
        "ex = relax.vm.build(MyModuleMixture, target=\"llvm\")\n",
        "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
        "\n",
        "nd_res = vm[\"main\"](data_nd,\n",
        "                    nd_params[\"w0\"],\n",
        "                    nd_params[\"b0\"],\n",
        "                    nd_params[\"w1\"],\n",
        "                    nd_params[\"b1\"])\n",
        "\n",
        "pred_kind = np.argmax(nd_res.numpy(), axis=1)\n",
        "print(\"MyModuleMixture Prediction:\", class_names[pred_kind[0]])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BoZGZXgsTgd2"
      },
      "source": [
        "## 参数绑定到IRModule\n",
        "之前的示例中，通过显式传递参数来构造主函数。将参数绑定为附加到IRModule的常量会降低API的复杂度。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HXBTZsrQUDYA",
        "outputId": "0aedb71e-a60e-43cb-d827-c029103acfc4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "@tvm.script.ir_module\n",
            "class Module:\n",
            "    @R.function\n",
            "    def main(x: Tensor((1, 784), \"float32\")) -> Tensor(None, \"float32\", ndim = 2):\n",
            "        # block 0\n",
            "        with R.dataflow():\n",
            "            lv0 = R.call_tir(linear0, (x, meta[relay.Constant][0], meta[relay.Constant][1]), (1, 128), dtype=\"float32\")\n",
            "            lv1 = R.call_tir(\"env.relu\", (lv0,), (1, 128), dtype=\"float32\")\n",
            "            out = R.call_tir(\"env.linear\", (lv1, meta[relay.Constant][2], meta[relay.Constant][3]), (1, 10), dtype=\"float32\")\n",
            "            R.output(out)\n",
            "        return out\n",
            "    \n",
            "    @T.prim_func\n",
            "    def linear0(X: T.Buffer[(1, 784), \"float32\"], W: T.Buffer[(128, 784), \"float32\"], B: T.Buffer[128, \"float32\"], Z: T.Buffer[(1, 128), \"float32\"]) -> None:\n",
            "        # function attr dict\n",
            "        T.func_attr({\"global_symbol\": \"linear0\", \"tir.noalias\": True})\n",
            "        # body\n",
            "        # with T.block(\"root\")\n",
            "        Y = T.alloc_buffer([1, 128], dtype=\"float32\")\n",
            "        for i, j, k in T.grid(1, 128, 784):\n",
            "            with T.block(\"Y\"):\n",
            "                vi, vj, vk = T.axis.remap(\"SSR\", [i, j, k])\n",
            "                T.reads(X[vi, vk], W[vj, vk])\n",
            "                T.writes(Y[vi, vj])\n",
            "                with T.init():\n",
            "                    Y[vi, vj] = T.float32(0)\n",
            "                Y[vi, vj] = Y[vi, vj] + X[vi, vk] * W[vj, vk]\n",
            "        for i, j in T.grid(1, 128):\n",
            "            with T.block(\"Z\"):\n",
            "                vi, vj = T.axis.remap(\"SS\", [i, j])\n",
            "                T.reads(Y[vi, vj], B[vj])\n",
            "                T.writes(Z[vi, vj])\n",
            "                Z[vi, vj] = Y[vi, vj] + B[vj]\n",
            "    \n"
          ]
        }
      ],
      "source": [
        "# 以下代码通过将参数名称与 nd_params 中的键匹配来创建绑定。\n",
        "MyModuleWithParams = relax.transform.BindParams(\"main\", nd_params)(MyModuleMixture)\n",
        "print(MyModuleWithParams.script())\n",
        "# IPython.display.Code(MyModuleWithParams.script(), language=\"python\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FG8yfCiJUSkz"
      },
      "source": [
        "在上面的脚本中，meta[relay.Constant][0] （译者注：目前 Relax 的常量表达依然继承自 Relay ，未来该 API 可能会更改） 对应于一个存储常量的隐式字典（它没有显示为脚本的一部分，但仍然是 IRModule 的一部分）。 如果我们构建转换后的 IRModule，我们现在可以通过传入输入数据来调用该函数。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pcVIc-CIUY0c",
        "outputId": "58ed71ce-acfb-41d4-c01e-5ab766ff7801"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "MyModuleWithParams Prediction: Sneaker\n"
          ]
        }
      ],
      "source": [
        "ex = relax.vm.build(MyModuleWithParams, target=\"llvm\")\n",
        "vm = relax.VirtualMachine(ex, tvm.cpu())\n",
        "\n",
        "nd_res = vm[\"main\"](data_nd)\n",
        "\n",
        "pred_kind = np.argmax(nd_res.numpy(), axis=1)\n",
        "print(\"MyModuleWithParams Prediction:\", class_names[pred_kind[0]])"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "GPS1VfL4mHpc",
        "o6hl9dZuJ4Pc"
      ],
      "name": "Untitled0.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3.8.10 ('base': conda)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.8.10"
    },
    "vscode": {
      "interpreter": {
        "hash": "2f394aca7ca06fed1e6064aef884364492d7cdda3614a461e02e6407fc40ba69"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}