tinynn_pytorch_to_tflite_int8.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cQQJJhk0ofMu",
        "outputId": "07ba3f62-786d-41ba-fbd8-9bdb340260a8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting git+https://github.com/alibaba/TinyNeuralNetwork.git\n",
            "  Cloning https://github.com/alibaba/TinyNeuralNetwork.git to /tmp/pip-req-build-lhes23rn\n",
            "  Running command git clone -q https://github.com/alibaba/TinyNeuralNetwork.git /tmp/pip-req-build-lhes23rn\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "    Preparing wheel metadata ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: ruamel.yaml>=0.16.12 in /usr/local/lib/python3.7/dist-packages (from TinyNeuralNetwork==0.1.0.20220929231907+a521163753e4b201a4a31d3cb82e5aa1fe328d93) (0.17.21)\n",
            "Requirement already satisfied: flatbuffers>=1.12 in /usr/local/lib/python3.7/dist-packages (from TinyNeuralNetwork==0.1.0.20220929231907+a521163753e4b201a4a31d3cb82e5aa1fe328d93) (2.0.7)\n",
            "Requirement already satisfied: igraph>=0.9 in /usr/local/lib/python3.7/dist-packages (from TinyNeuralNetwork==0.1.0.20220929231907+a521163753e4b201a4a31d3cb82e5aa1fe328d93) (0.10.1)\n",
            "Requirement already satisfied: PyYAML>=5.3.1 in /usr/local/lib/python3.7/dist-packages (from TinyNeuralNetwork==0.1.0.20220929231907+a521163753e4b201a4a31d3cb82e5aa1fe328d93) (6.0)\n",
            "Requirement already satisfied: numpy>=1.18.5 in /usr/local/lib/python3.7/dist-packages (from TinyNeuralNetwork==0.1.0.20220929231907+a521163753e4b201a4a31d3cb82e5aa1fe328d93) (1.21.6)\n",
            "Requirement already satisfied: texttable>=1.6.2 in /usr/local/lib/python3.7/dist-packages (from igraph>=0.9->TinyNeuralNetwork==0.1.0.20220929231907+a521163753e4b201a4a31d3cb82e5aa1fe328d93) (1.6.4)\n",
            "Requirement already satisfied: ruamel.yaml.clib>=0.2.6 in /usr/local/lib/python3.7/dist-packages (from ruamel.yaml>=0.16.12->TinyNeuralNetwork==0.1.0.20220929231907+a521163753e4b201a4a31d3cb82e5aa1fe328d93) (0.2.6)\n",
            "Building wheels for collected packages: TinyNeuralNetwork\n",
            "  Building wheel for TinyNeuralNetwork (PEP 517) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for TinyNeuralNetwork: filename=TinyNeuralNetwork-0.1.0.20220929231908+a521163753e4b201a4a31d3cb82e5aa1fe328d93-py3-none-any.whl size=352164 sha256=4dc00262561963e130e096be1d3dccf54d4c7164f601140b0974d12545c0bc9b\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-rkf9udsc/wheels/64/6b/f5/9cf69de054ba0de53e572bc2f13988a664f3dc9623e6f0825d\n",
            "Successfully built TinyNeuralNetwork\n",
            "Installing collected packages: TinyNeuralNetwork\n",
            "  Attempting uninstall: TinyNeuralNetwork\n",
            "    Found existing installation: TinyNeuralNetwork 0.1.0.20220929231341+a521163753e4b201a4a31d3cb82e5aa1fe328d93\n",
            "    Uninstalling TinyNeuralNetwork-0.1.0.20220929231341+a521163753e4b201a4a31d3cb82e5aa1fe328d93:\n",
            "      Successfully uninstalled TinyNeuralNetwork-0.1.0.20220929231341+a521163753e4b201a4a31d3cb82e5aa1fe328d93\n",
            "Successfully installed TinyNeuralNetwork-0.1.0.20220929231908+a521163753e4b201a4a31d3cb82e5aa1fe328d93\n"
          ]
        }
      ],
      "source": [
        "!pip install git+https://github.com/alibaba/TinyNeuralNetwork.git"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install git+https://github.com/openai/whisper.git "
      ],
      "metadata": {
        "id": "zf4Er13nq7gg",
        "outputId": "bfa6e470-482d-4bec-d6a0-c149c905329b",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting git+https://github.com/openai/whisper.git\n",
            "  Cloning https://github.com/openai/whisper.git to /tmp/pip-req-build-fgx6fwr7\n",
            "  Running command git clone -q https://github.com/openai/whisper.git /tmp/pip-req-build-fgx6fwr7\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from whisper==1.0) (1.21.6)\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (from whisper==1.0) (1.12.1+cu113)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from whisper==1.0) (4.64.1)\n",
            "Requirement already satisfied: more-itertools in /usr/local/lib/python3.7/dist-packages (from whisper==1.0) (8.14.0)\n",
            "Requirement already satisfied: transformers>=4.19.0 in /usr/local/lib/python3.7/dist-packages (from whisper==1.0) (4.22.2)\n",
            "Requirement already satisfied: ffmpeg-python==0.2.0 in /usr/local/lib/python3.7/dist-packages (from whisper==1.0) (0.2.0)\n",
            "Requirement already satisfied: future in /usr/local/lib/python3.7/dist-packages (from ffmpeg-python==0.2.0->whisper==1.0) (0.16.0)\n",
            "Requirement already satisfied: tokenizers!=0.11.3,<0.13,>=0.11.1 in /usr/local/lib/python3.7/dist-packages (from transformers>=4.19.0->whisper==1.0) (0.12.1)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.9.0 in /usr/local/lib/python3.7/dist-packages (from transformers>=4.19.0->whisper==1.0) (0.10.0)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from transformers>=4.19.0->whisper==1.0) (4.12.0)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.7/dist-packages (from transformers>=4.19.0->whisper==1.0) (2022.6.2)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from transformers>=4.19.0->whisper==1.0) (2.23.0)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.7/dist-packages (from transformers>=4.19.0->whisper==1.0) (6.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from transformers>=4.19.0->whisper==1.0) (3.8.0)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.7/dist-packages (from transformers>=4.19.0->whisper==1.0) (21.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.7/dist-packages (from huggingface-hub<1.0,>=0.9.0->transformers>=4.19.0->whisper==1.0) (4.1.1)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging>=20.0->transformers>=4.19.0->whisper==1.0) (3.0.9)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->transformers>=4.19.0->whisper==1.0) (3.8.1)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->transformers>=4.19.0->whisper==1.0) (3.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->transformers>=4.19.0->whisper==1.0) (2022.6.15)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->transformers>=4.19.0->whisper==1.0) (2.10)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->transformers>=4.19.0->whisper==1.0) (1.24.3)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/usefulsensors/openai-whisper.git"
      ],
      "metadata": {
        "id": "zN9Qqqg_CPcZ",
        "outputId": "a0614946-fdfe-4484-c027-7b263d62a797",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'openai-whisper'...\n",
            "remote: Enumerating objects: 33, done.\u001b[K\n",
            "remote: Counting objects: 100% (33/33), done.\u001b[K\n",
            "remote: Compressing objects: 100% (32/32), done.\u001b[K\n",
            "remote: Total 33 (delta 14), reused 0 (delta 0), pack-reused 0\u001b[K\n",
            "Unpacking objects: 100% (33/33), done.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Convert from pytorch to onnx"
      ],
      "metadata": {
        "id": "vAe5oHk1i0iI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import whisper\n",
        "import torch\n",
        "\n",
        "tiny_model = whisper.load_model(\"tiny\")\n",
        "torch.onnx.export(tiny_model.encoder, torch.randn(1,80,3000).to(\"cuda\"), \"/content/whisper-encoder.onnx\")\n",
        "torch.onnx.export(tiny_model.decoder, (torch.tensor([[50258]]).to(\"cuda\"), torch.randn(1,384,384).to(\"cuda\")), \"/content/whisper-decoder_main.onnx\")\n",
        "torch.onnx.export(tiny_model.decoder, (torch.tensor([[50258, 50259, 50359]]).to(\"cuda\"), torch.randn(1, 384, 384).to(\"cuda\")), \"/content/whisper-decoder_language.onnx\")"
      ],
      "metadata": {
        "id": "DeDY0j0zitxx"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Convert from pytorch to tflite using Tinynn"
      ],
      "metadata": {
        "id": "yFpdUZS2ivnh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import whisper\n",
        "import torch\n",
        "\n",
        "import torch.nn as nn\n",
        "from tinynn.converter import TFLiteConverter\n",
        "from tinynn.graph.quantization.quantizer import PostQuantizer\n",
        "from tinynn.graph.tracer import model_tracer\n",
        "from tinynn.util.train_util import DLContext, get_device, train\n",
        "\n",
        "model = whisper.load_model(\"tiny\")\n",
        "model.eval()\n",
        "audio = whisper.load_audio(\"/content/openai-whisper/filename.mp3\")\n",
        "audio = whisper.pad_or_trim(audio)\n",
        "print(audio.shape)\n",
        "# make log-Mel spectrogram and move to the same device as the model\n",
        "mel = whisper.log_mel_spectrogram(audio).to(model.device)\n",
        "print(mel.shape)\n",
        "quantizer = PostQuantizer(model, mel, work_dir='out', config={'asymmetric': True, 'per_tensor': False})\n",
        "#qat_model = quantizer.quantize()\n"
      ],
      "metadata": {
        "id": "ruITGopwrCz2",
        "outputId": "495d8c2a-31ac-42e2-db65-7eee1b859735",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "(480000,)\n",
            "torch.Size([80, 3000])\n"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "tinynn_pytorch_to_tflite_int8.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}