PyTorch_To_ONNX_Model_Converter_Google_Colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "private_outputs": true,
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# PyTorch To ONNX Model Converter Google Colab (ver. 1.0)\n",
        "\n",
        "***\n",
        "\n",
        "Powered by tegridy-tools: https://github.com/asigalov61/tegridy-tools\n",
        "\n",
        "***\n",
        "\n",
        "WARNING: This complete implementation is a functioning model of the Artificial Intelligence. Please excercise great humility, care, and respect. https://www.nscai.gov/\n",
        "\n",
        "***\n",
        "\n",
        "#### Project Los Angeles\n",
        "\n",
        "#### Tegridy Code 2023\n",
        "\n",
        "***"
      ],
      "metadata": {
        "id": "Zw8vqR-yFqDq"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# NVIDIA GPU Check"
      ],
      "metadata": {
        "id": "LgUeL9cx-qoD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!nvidia-smi"
      ],
      "metadata": {
        "id": "E4RuZ9X4TdLB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Setup Environment\n",
        "\n",
        "## Run cells below one at a time and make sure to restart the colab runtime/re-run each cell again after the run to make sure that there are no residual install errors or conflicts"
      ],
      "metadata": {
        "id": "OQMiTGp8-u7q"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip uninstall torch"
      ],
      "metadata": {
        "id": "D4vxN5zSYQl0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121"
      ],
      "metadata": {
        "id": "9J75Ps7uYHGy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install numpy -U"
      ],
      "metadata": {
        "id": "xXuD9k0hXNrz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install transformers -U\n",
        "!pip install onnx -U\n",
        "!pip install onnxruntime -U\n",
        "!pip install onnxruntime-gpu -U\n",
        "!pip install einops -U\n",
        "!pip install huggingface_hub -U"
      ],
      "metadata": {
        "id": "iGSolqJk4_QT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prep and Load Your (Custom) PyTorch Model/Architecture\n",
        "\n",
        "## Below we are going to use Allegro Music Transformer PyTorch Music AI model/custom architecture as an example"
      ],
      "metadata": {
        "id": "T6ZiG0pmKOew"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Download the PyTorch model"
      ],
      "metadata": {
        "id": "LQ97Nhbk_tVl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from huggingface_hub import hf_hub_download\n",
        "\n",
        "hf_hub_download(repo_id='asigalov61/Allegro-Music-Transformer',\n",
        "                filename='Allegro_Music_Transformer_Small_Trained_Model_56000_steps_0.9399_loss_0.7374_acc.pth',\n",
        "                local_dir='/content/',\n",
        "                local_dir_use_symlinks=False)"
      ],
      "metadata": {
        "id": "b1AT3L0XBmRm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prep custom PyTorch model environment/architecture"
      ],
      "metadata": {
        "id": "_13e8yDRAy0w"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/asigalov61/tegridy-tools\n",
        "\n",
        "%cd /content/tegridy-tools/tegridy-tools\n",
        "\n",
        "import TMIDIX\n",
        "\n",
        "%cd /content/tegridy-tools/tegridy-tools/X-Transformer\n",
        "\n",
        "from x_transformer import *\n",
        "\n",
        "%cd /content/"
      ],
      "metadata": {
        "id": "FjQzg9I2AOQH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Load custom PyTorch model"
      ],
      "metadata": {
        "id": "7M8ylIP9BLkI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "\n",
        "device = torch.device('cuda:0')\n",
        "\n",
        "SEQ_LEN = 2048\n",
        "\n",
        "# instantiate the model\n",
        "\n",
        "model = TransformerWrapper(\n",
        "    num_tokens = 3088,\n",
        "    max_seq_len = SEQ_LEN,\n",
        "    attn_layers = Decoder(dim = 1024, depth = 32, heads = 8)\n",
        ").cuda()\n",
        "\n",
        "model = AutoregressiveWrapper(model).cuda() #, mask_prob = 0.15)\n",
        "model = torch.nn.DataParallel(model).cuda()\n",
        "model.cuda()\n",
        "model.load_state_dict(torch.load('/content/Allegro_Music_Transformer_Small_Trained_Model_56000_steps_0.9399_loss_0.7374_acc.pth'))\n",
        "model.cuda()\n",
        "\n",
        "\n",
        "\n",
        "model.eval()\n",
        "model.cuda()\n",
        "\n",
        "model.to(device=device)\n",
        "print('Done!')"
      ],
      "metadata": {
        "id": "LuyUsaoxcvXi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Generate seq_len sized output from the PyTorch model to use with PyTorch ONNX exporter. Save it for later use as well.\n",
        "\n",
        "## Please note that generating seq_len output may take up-to 10 minutes on T4 GPU so please be patient :)"
      ],
      "metadata": {
        "id": "Z9tmnqUcDK4i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "x = torch.LongTensor([[2816]]).cuda()\n",
        "\n",
        "torch_out = model.module.generate(x, 2048)"
      ],
      "metadata": {
        "id": "s1-hq885wMwW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch_out[0].tolist()"
      ],
      "metadata": {
        "id": "SFyxAmZ76vhW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Save the generated output from the source model"
      ],
      "metadata": {
        "id": "SNzLJE8AFvl_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with open('torch_out.txt', 'w') as file:\n",
        "   file.write('\\n'.join(str(to) for to in torch_out[0].tolist()))"
      ],
      "metadata": {
        "id": "sSa7uwu5JkVV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Load and prep the generated output from the source model for use with ONNX exporter"
      ],
      "metadata": {
        "id": "LZ8vVWoUFz9v"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with open(\"torch_out.txt\", 'r') as file:\n",
        "  t_out = file.readlines()\n",
        "\n",
        "torch_out1 = []\n",
        "for t in t_out:\n",
        "  torch_out1.append(int(t))"
      ],
      "metadata": {
        "id": "1nOW8wwjKXpn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch_out2 = torch.LongTensor([torch_out1]).to(device=device)"
      ],
      "metadata": {
        "id": "N0kxK86sMcP8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch_out2"
      ],
      "metadata": {
        "id": "C2uF_KL9MjBK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Start the ONNX export/conversion to ONNX format!\n",
        "\n",
        "## You can ignore grey warning messages from the exporter!"
      ],
      "metadata": {
        "id": "izcm3G0iGQp2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "\n",
        "\n",
        "def export_onnx(mod, model_inputs, path):\n",
        "\n",
        "    torch.onnx.export(mod,  # model being run\n",
        "                      model_inputs,  # model input (or a tuple for multiple inputs)\n",
        "                      path,  # where to save the model (can be a file or file-like object)\n",
        "                      export_params=True,  # store the trained parameter weights inside the model file\n",
        "                      opset_version=14,  # the ONNX version to export the model to\n",
        "                      do_constant_folding=True,  # whether to execute constant folding for optimization\n",
        "                      verbose=True,\n",
        "                      input_names=['input'],\n",
        "                      output_names=['output'],\n",
        "                      dynamic_axes={ # dict value: manually named axes\n",
        "                                    \"input\": {1: \"seq_len\"},\n",
        "                                    # list value: automatic names\n",
        "                                    \"output\": [0],\n",
        "                                }\n",
        "                      )\n",
        "\n",
        "    print('finished exporting onnx')\n",
        "\n",
        "export_onnx(model.module.net.to(device=device), torch_out2, path='/content/test.bin')"
      ],
      "metadata": {
        "id": "0SpaI6eziJYS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Test the resulting ONNX model (graph)"
      ],
      "metadata": {
        "id": "9FFOpndZG3m8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import onnx\n",
        "\n",
        "# Load the ONNX model\n",
        "omodel = onnx.load(\"/content/test.bin\")\n",
        "\n",
        "# Check that the model is well formed\n",
        "onnx.checker.check_model(omodel)\n",
        "\n",
        "# Print a human readable representation of the graph\n",
        "print(onnx.helper.printable_graph(omodel.graph))"
      ],
      "metadata": {
        "id": "gOvQA15l61G0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Check IO names, shapes, and types"
      ],
      "metadata": {
        "id": "B6CCKBBPHAXL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import onnxruntime as ort\n",
        "\n",
        "session = ort.InferenceSession('/content/test.bin', None, providers=['CUDAExecutionProvider'])\n",
        "input_name = session.get_inputs()[0].name\n",
        "output_name = session.get_outputs()[0].name\n",
        "print(input_name)\n",
        "print(output_name)"
      ],
      "metadata": {
        "id": "IMK9lb5q3spT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "input_name = session.get_inputs()[0].name\n",
        "print(\"input name\", input_name)\n",
        "input_shape = session.get_inputs()[0].shape\n",
        "print(\"input shape\", input_shape)\n",
        "input_type = session.get_inputs()[0].type\n",
        "print(\"input type\", input_type)\n",
        "\n",
        "output_name = session.get_outputs()[0].name\n",
        "print(\"output name\", output_name)\n",
        "output_shape = session.get_outputs()[0].shape\n",
        "print(\"output shape\", output_shape)\n",
        "output_type = session.get_outputs()[0].type\n",
        "print(\"output type\", output_type)"
      ],
      "metadata": {
        "id": "c4an_abnB4WT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Now we are going to test the resulting ONNX model for inference!"
      ],
      "metadata": {
        "id": "pHAV8Z20HXED"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Load the sample model input here to prime the model with"
      ],
      "metadata": {
        "id": "AGVjGr4QHkYD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "with open(\"torch_out.txt\", 'r') as file:\n",
        "  t_out = file.readlines()\n",
        "\n",
        "torch_out1 = []\n",
        "for t in t_out:\n",
        "  torch_out1.append(int(t))"
      ],
      "metadata": {
        "id": "s199ECPse7E1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Initialize the ONNX model!"
      ],
      "metadata": {
        "id": "VwY1b4cUH0LC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import onnxruntime\n",
        "\n",
        "session = onnxruntime.InferenceSession('/content/test.bin', providers=['CUDAExecutionProvider'])"
      ],
      "metadata": {
        "id": "XQIyIs-SCbXF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Below we are going to init a simple autoregressive wrapper function to eval the model and the output!"
      ],
      "metadata": {
        "id": "HSdEtGoWIH2O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn.functional as F\n",
        "\n",
        "#====================================================\n",
        "\n",
        "@torch.no_grad()\n",
        "def generate(\n",
        "            start_tokens,\n",
        "            seq_len,\n",
        "            max_seq_len = 2048,\n",
        "            temperature = 0.9,\n",
        "            verbose=True,\n",
        "            return_prime=False,\n",
        "            ):\n",
        "\n",
        "    out = torch.FloatTensor([start_tokens])\n",
        "\n",
        "    st = len(start_tokens)\n",
        "\n",
        "    if verbose:\n",
        "      print(\"Generating sequence of max length:\", seq_len)\n",
        "\n",
        "    for s in range(seq_len):\n",
        "        x = out[:, -max_seq_len:]\n",
        "\n",
        "        torch_in = x.tolist()[0]\n",
        "\n",
        "        logits = torch.FloatTensor(session.run(None, {'input': [torch_in]})[0])[:, -1]\n",
        "\n",
        "        filtered_logits = logits\n",
        "\n",
        "        probs = F.softmax(filtered_logits / temperature, dim=-1)\n",
        "\n",
        "        sample = torch.multinomial(probs, 1)\n",
        "\n",
        "        out = torch.cat((out, sample), dim=-1)\n",
        "\n",
        "        if verbose:\n",
        "          if s % 32 == 0:\n",
        "            print(s, '/', seq_len)\n",
        "\n",
        "    if return_prime:\n",
        "      return out[:, :]\n",
        "\n",
        "    else:\n",
        "      return out[:, st:]"
      ],
      "metadata": {
        "id": "Wf3F7JQGLMG6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Generation time!!!"
      ],
      "metadata": {
        "id": "LrxlQGiuIdVy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "melody_chords_f = generate([3087, 3073+1, 3075+1], 512)"
      ],
      "metadata": {
        "id": "E-EEKBIlQiCL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Convert the generated tokens/output to MIDI and enjoy! :)"
      ],
      "metadata": {
        "id": "grNwJmcxIpc-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import TMIDIX\n",
        "\n",
        "melody_chords_f = melody_chords_f.tolist()[0]\n",
        "\n",
        "print('=' * 70)\n",
        "print('Sample INTs', melody_chords_f[:12])\n",
        "print('=' * 70)\n",
        "\n",
        "if len(melody_chords_f) != 0:\n",
        "\n",
        "    song = melody_chords_f\n",
        "    song_f = []\n",
        "    time = 0\n",
        "    dur = 0\n",
        "    vel = 0\n",
        "    pitch = 0\n",
        "    channel = 0\n",
        "\n",
        "    for sss in song:\n",
        "\n",
        "      ss = int(sss)\n",
        "\n",
        "      if ss > 0 and ss < 256:\n",
        "\n",
        "          time += ss * 8\n",
        "\n",
        "      if ss >= 256 and ss < 1280:\n",
        "\n",
        "          dur = ((ss-256) // 8) * 32\n",
        "          vel = (((ss-256) % 8)+1) * 15\n",
        "\n",
        "      if ss >= 1280 and ss < 2816:\n",
        "          channel = (ss-1280) // 128\n",
        "          pitch = (ss-1280) % 128\n",
        "\n",
        "          song_f.append(['note', time, dur, channel, pitch, vel ])\n",
        "\n",
        "    detailed_stats = TMIDIX.Tegridy_SONG_to_MIDI_Converter(song_f,\n",
        "                                                        output_signature = 'Allegro Music Transformer',\n",
        "                                                        output_file_name = '/content/Allegro-Music-Transformer-Music-Composition',\n",
        "                                                        track_name='Project Los Angeles',\n",
        "                                                        list_of_MIDI_patches=[0, 24, 32, 40, 42, 46, 56, 71, 73, 0, 53, 19, 0, 0, 0, 0],\n",
        "                                                        number_of_ticks_per_quarter=500)\n",
        "    print('=' * 70)"
      ],
      "metadata": {
        "id": "1iaUFYAndcAY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Congrats! You did it! :)"
      ],
      "metadata": {
        "id": "11_v4fMjI0oe"
      }
    }
  ]
}