Export_CLIP_to_ONNX_tflite_tfjs_tf_saved_model.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Export CLIP to ONNX/tflite/tfjs/tf saved model.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "AtWchYtU0Dtv"
      },
      "source": [
        "# IMPORTANT: Make sure you're using a GPU runtime!"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h_zSu-EKxlBP"
      },
      "source": [
        "# Based on this notebook: https://colab.research.google.com/github/openai/clip/blob/master/notebooks/Interacting_with_CLIP.ipynb"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ur4qmQUIBxwe"
      },
      "source": [
        "!git clone https://github.com/openai/CLIP\n",
        "%cd CLIP"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Rbzq8REAy73u"
      },
      "source": [
        "# This is SUPER hacky because I don't know a better way (that's quick). Basically the vision model is ready to export as-is, like this:\n",
        "#   torch.onnx.export(model.vision, ...)\n",
        "# but the text model has a couple of pre-processing steps (like converting tokens to embeddings), and I'd like to have all that\n",
        "# processing contained within the onnx file for the text encoder. The `torch.onnx.export` function seems to only be able to\n",
        "# take a *model* as an input, and not a function (like `model.encode_text`), so I'm hackily renaming `model.encode_text` to\n",
        "# `model.forward` so that I can then write:\n",
        "#   torch.onnx.export(model, ...)\n",
        "# to export the text encoder. I'm sure there's a much better way to do this. If this stops working, note that\n",
        "# it was working at the following commit hash, so you can clone this to get it working: https://github.com/openai/CLIP/tree/573315e83f07b53a61ff5098757e8fc885f1703e\n",
        "!sed -i -e 's/def forward(self, image, text):/def old_forward(self, image, text):/g' ./clip/model.py\n",
        "!sed -i -e 's/def encode_text(self, text):/def forward(self, text):/g' ./clip/model.py"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0BpdJkdBssk9"
      },
      "source": [
        "! pip install ftfy regex tqdm"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uLFS29hnhlY4"
      },
      "source": [
        "import numpy as np\n",
        "import torch\n",
        "import clip\n",
        "\n",
        "clip.available_models()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IBRVTY9lbGm8"
      },
      "source": [
        "model, preprocess = clip.load(\"ViT-B/32\")\n",
        "model.cuda().eval()\n",
        "input_resolution = model.visual.input_resolution\n",
        "context_length = model.context_length\n",
        "vocab_size = model.vocab_size\n",
        "\n",
        "print(\"Model parameters:\", f\"{np.sum([int(np.prod(p.shape)) for p in model.parameters()]):,}\")\n",
        "print(\"Input resolution:\", input_resolution)\n",
        "print(\"Context length:\", context_length)\n",
        "print(\"Vocab size:\", vocab_size)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qGom156-i2kL"
      },
      "source": [
        "clip.tokenize(\"Hello World!\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tMc1AXzBlhzm"
      },
      "source": [
        "import os\n",
        "import skimage\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "\n",
        "from collections import OrderedDict\n",
        "import torch\n",
        "\n",
        "# images in skimage to use and their textual descriptions\n",
        "descriptions = {\n",
        "    \"astronaut\": \"a portrait of an astronaut with the American flag\",\n",
        "}"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NSSrLY185jSf"
      },
      "source": [
        "original_images = []\n",
        "images = []\n",
        "texts = []\n",
        "\n",
        "for filename in [filename for filename in os.listdir(skimage.data_dir) if filename.endswith(\".png\") or filename.endswith(\".jpg\")]:\n",
        "    name = os.path.splitext(filename)[0]\n",
        "    if name not in descriptions:\n",
        "        continue\n",
        "\n",
        "    image = Image.open(os.path.join(skimage.data_dir, filename)).convert(\"RGB\")\n",
        "    original_images.append(image)\n",
        "    images.append(preprocess(image))\n",
        "    texts.append(descriptions[name])"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HBgCanxi8JKw"
      },
      "source": [
        "image_input = torch.tensor(np.stack(images)).half().cuda()\n",
        "text_tokens = clip.tokenize([\"This is \" + desc for desc in texts]).cuda()"
      ],
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model.visual(image_input)[0] # astronaut pic embedding"
      ],
      "metadata": {
        "id": "g0o8mDN6wq_L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model(text_tokens)[0] # astronaut text embedding"
      ],
      "metadata": {
        "id": "qEPHMWwN0Puv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kDmmi0vMI9WY"
      },
      "source": [
        "torch.onnx.export(model, text_tokens, \"clip-text-vit-32.onnx\", export_params=True, opset_version=12, do_constant_folding=True, input_names = ['input'], output_names = ['output'], dynamic_axes={'input' : {0 : 'batch_size'}, 'output' : {0 : 'batch_size'}})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BLSGVjueonP0"
      },
      "source": [
        "torch.onnx.export(model.visual, image_input, \"clip-image-vit-32.onnx\", export_params=True, opset_version=12, do_constant_folding=True, input_names = ['input'], output_names = ['output'], dynamic_axes={'input' : {0 : 'batch_size'}, 'output' : {0 : 'batch_size'}})"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X0I6iPCOxB9M"
      },
      "source": [
        "# use this option in the above torch.onnx.export calls if you get a \"Unable to cast from non-held to held instance (T& to Holder<T>)\" error:\n",
        "#   operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OhCoVnSo2XFr"
      },
      "source": [
        "# The onnx model files are now in the /content/CLIP directory."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Attempt at quantizing model to uint8 (doesn't seem to work? no errors, but onnx file is same size)\n",
        "# Reference: https://github.com/minimaxir/imgbeddings/blob/36fb4d7ac6b82694d109cef6f887d4cb9c49da0f/imgbeddings/models.py#L94\n",
        "# Here's the model the above code generates: https://huggingface.co/minimaxir/imgbeddings/blob/main/patch32_v1.onnx\n",
        "# Here's a demo of the above ONNX model with ORT Web: https://jsbin.com/nupehazaju/edit?html,output  <-- seems to work, but this model doesn't have the projection head that squashes 768 vec to 512 elements (so can be compared to text embeddings of same length)\n",
        "!pip install onnxruntime\n",
        "!pip install onnx\n",
        "from onnxruntime.quantization import quantize_dynamic, QuantType\n",
        "quantize_dynamic(\"clip-image-vit-32.onnx\", \"clip-image-vit-32-uint8.onnx\", weight_type=QuantType.QUInt8, extra_options={\"MatMulConstBOnly\":False}) # I added the MatMulConstBOnly as a guess due to warnings that it outputs without it"
      ],
      "metadata": {
        "id": "24LcAdP2doTx"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "44jzFoZzxPrf"
      },
      "source": [
        "# The code below is for converting to tflite, tfjs and tf saved model:"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r2VoXSsyyFu-"
      },
      "source": [
        "!pip install git+https://github.com/onnx/onnx-tensorflow.git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n0axzSah0_h4"
      },
      "source": [
        "!onnx-tf convert -i clip-image-vit-32.onnx -o clip-image-vit-32-tf\n",
        "!onnx-tf convert -i clip-text-vit-32.onnx -o clip-text-vit-32-tf"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2kDc0sPILbQu"
      },
      "source": [
        "!pip install tensorflowjs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WXFWVZACLUR8"
      },
      "source": [
        "!tensorflowjs_converter --input_format tf_saved_model ./clip-image-vit-32-tf ./clip-image-vit-32-tfjs\n",
        "!tensorflowjs_converter --input_format tf_saved_model ./clip-text-vit-32-tf ./clip-text-vit-32-tfjs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g1Ub_dsaKqO8"
      },
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "# image encoder:\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model(\"./clip-image-vit-32-tf\")\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS] # This line is needed because: https://github.com/tensorflow/tfjs/issues/5844\n",
        "tflite_model = converter.convert()\n",
        "with open('clip-image-vit-32.tflite', 'wb') as f:\n",
        "  f.write(tflite_model)\n",
        "\n",
        "# text encoder:\n",
        "converter = tf.lite.TFLiteConverter.from_saved_model(\"./clip-text-vit-32-tf\")\n",
        "converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS] # This line is needed because: https://github.com/tensorflow/tfjs/issues/5844\n",
        "tflite_model = converter.convert()\n",
        "with open('clip-text-vit-32.tflite', 'wb') as f:\n",
        "  f.write(tflite_model)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}