quantization.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNnkRvLASzLOuVEZWfnPl2w",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Rumeysakeskin/ASR-Quantization/blob/main/quantization.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G_wSnCIqP0UH"
      },
      "outputs": [],
      "source": [
        "# If you're using Google Colab and not running locally, run this cell.\n",
        "## Install dependencies\n",
        "!pip install wget\n",
        "!apt-get install sox libsndfile1 ffmpeg\n",
        "!pip install text-unidecode\n",
        "!pip install matplotlib>=3.3.2\n",
        "\n",
        "## Install NeMo\n",
        "BRANCH = 'main'\n",
        "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n",
        "!apt-get update && apt-get install -y libsndfile1 ffmpeg\n",
        "!pip install Cython tensorflow==2.11.0 Pygments==2.6.1 pynini==2.1.5 nemo_toolkit[all]"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import nemo.collections.asr as nemo_asr\n",
        "from ruamel.yaml import YAML\n",
        "import pytorch_lightning as pl\n",
        "from omegaconf import DictConfig, OmegaConf, open_dict\n",
        "import copy\n",
        "from pytorch_lightning.callbacks import ModelCheckpoint\n",
        "import os\n",
        "import torch.quantization as quant\n",
        "\n",
        "if not os.path.exists(\"configs/config.yaml\"):\n",
        "   !wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/$BRANCH/examples/asr/conf/config.yaml"
      ],
      "metadata": {
        "id": "jnnLqh6LWHKc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "FINE TUNING"
      ],
      "metadata": {
        "id": "VOL-Awq1WMgz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def training_quartznet():\n",
        "\n",
        "    EPOCHS = 100\n",
        "   \n",
        "    config_path = \"configs/config.yaml\"\n",
        "\n",
        "    yaml = YAML(typ='safe')\n",
        "    with open(config_path) as f:\n",
        "        params = yaml.load(f)\n",
        "\n",
        "    params['model']['train_ds']['manifest_filepath'] = \"data/train_manifest.jsonl\"\n",
        "    params['model']['validation_ds']['manifest_filepath'] = \"data/val_manifest.jsonl\"\n",
        "\n",
        "\n",
        "    first_asr_model = nemo_asr.models.EncDecCTCModel.from_pretrained(\"QuartzNet15x5Base-En\")\n",
        "\n",
        "    # Insert the QuantStub() before the first layer of the model\n",
        "    first_asr_model.quant = torch.quantization.QuantStub()\n",
        "    first_asr_model.encoder.quant = torch.quantization.QuantStub()\n",
        "\n",
        "    # Insert a DeQuantStub() at the end of the model\n",
        "    first_asr_model.dequant = torch.quantization.DeQuantStub()\n",
        "    first_asr_model.decoder.dequant = torch.quantization.DeQuantStub()\n",
        "\n",
        "    \n",
        "    first_asr_model.change_vocabulary(\n",
        "        new_vocabulary=[\" \", \"a\", \"b\", \"c\", \"ç\", \"d\", \"e\", \"f\", \"g\", \"ğ\", \"h\", \"ı\", \"i\", \"j\", \"k\", \"l\", \"m\",\n",
        "                        \"n\", \"o\", \"ö\", \"p\", \"q\", \"r\", \"s\", \"ş\", \"t\", \"u\", \"ü\", \"v\", \"w\", \"x\", \"y\", \"z\", \"'\"])\n",
        "\n",
        "    new_opt = copy.deepcopy(params['model']['optim'])\n",
        "\n",
        "    new_opt['lr'] = 0.001\n",
        "    # Point to the data we'll use for fine-tuning as the training set\n",
        "    first_asr_model.setup_training_data(train_data_config=params['model']['train_ds'])\n",
        "    # Point to the new validation data for fine-tuning\n",
        "    first_asr_model.setup_validation_data(val_data_config=params['model']['validation_ds'])\n",
        "    # assign optimizer config\n",
        "    first_asr_model.setup_optimization(optim_config=DictConfig(new_opt))\n",
        "\n",
        "\n",
        "    # used for saving models\n",
        "    save_path = os.path.join(os.getcwd(),\"Quartznet15x5_models\")\n",
        "    checkpoint_callback = ModelCheckpoint(\n",
        "        dirpath=save_path,\n",
        "        save_top_k= -1,\n",
        "        verbose=True,\n",
        "        monitor='val_loss',\n",
        "        mode='min',\n",
        "    )\n",
        "    \n",
        "    trainer = pl.Trainer(gpus=N_GPUS, accelerator='ddp',num_nodes=N_DEVICES,\n",
        "                         max_epochs=EPOCHS, amp_level='O1',\n",
        "                         log_every_n_steps=150,\n",
        "                         val_check_interval=1.0, checkpoint_callback=checkpoint_callback)\n",
        "\n",
        "    first_asr_model.set_trainer(trainer)\n",
        "\n",
        "    trainer.fit(first_asr_model)\n",
        "\n",
        "if __name__ == '__main__':\n",
        "    training_quartznet()"
      ],
      "metadata": {
        "id": "0ctFggRwQOdr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "SAVE AND LOAD QUANTIZED PYTORCH MODEL"
      ],
      "metadata": {
        "id": "b9Uhxwu6Wk26"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# attach a global qconfig, which contains information about what kind\n",
        "# of observers to attach. Use 'fbgemm' for server inference and\n",
        "# 'qnnpack' for mobile inference. Other quantization configurations such\n",
        "# as selecting symmetric or assymetric quantization and MinMax or L2Norm\n",
        "# calibration techniques can be specified here.\n",
        "first_asr_model.qconfig = torch.quantization.get_default_qconfig('qnnpack')\n",
        "# Prepare the model for static quantization. This inserts observers in\n",
        "# the model that will observe activation tensors during calibration.\n",
        "torch.quantization.prepare(first_asr_model, inplace=True)\n",
        "# Convert the observed model to a quantized model. This does several things:\n",
        "# quantizes the weights, computes and stores the scale and bias value to be\n",
        "# used with each activation tensor, and replaces key operators with quantized\n",
        "# implementations.\n",
        "torch.quantization.convert(first_asr_model, inplace=True)\n",
        "\n",
        "model_filepath = os.path.join(\"/Quartznet15x5_models\", \"quantized_model.ckpt\")\n",
        "# print(\"first_asr_model.state_dict() name: \",first_asr_model.state_dict())\n",
        "torch.save(first_asr_model.state_dict(), model_filepath)"
      ],
      "metadata": {
        "id": "B5lJkDQFQ7_g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "INFERENCE WITH QUANTIZED MODEL"
      ],
      "metadata": {
        "id": "d4B-U4ikW6qW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model_to_load = os.path.join(\"/Quartznet15x5_models/quantized_model.ckpt\")\n",
        "config_path = \"configs/config.yaml\"\n",
        "\n",
        "yaml = YAML(typ='safe')\n",
        "with open(config_path) as f:\n",
        "    params = yaml.load(f)\n",
        "\n",
        "params['model']['train_ds']['manifest_filepath'] = \"data/train_manifest.jsonl\"\n",
        "params['model']['validation_ds']['manifest_filepath'] = \"data/val_manifest.jsonl\"\n",
        "\n",
        "first_asr_model = nemo_asr.models.EncDecCTCModel(cfg=DictConfig(params['model']))\n",
        "checkpoint = torch.load(model_to_load)\n",
        "\n",
        "for key in list(checkpoint.keys()):\n",
        "    if 'module.encoder.encoder' in key:\n",
        "        checkpoint[key.replace('module.', '')] = checkpoint.pop(key)\n",
        "for key in list(checkpoint.keys()):\n",
        "    if '.qconfig.' in key:\n",
        "        checkpoint.pop(key)\n",
        "\n",
        "# Dequantize the tensor before loading it\n",
        "for key, value in checkpoint.items():\n",
        "    if value.is_quantized:\n",
        "        checkpoint[key] = value.dequantize()\n",
        "\n",
        "first_asr_model.load_state_dict(checkpoint, strict=False)\n",
        "# model must be set to eval mode for static quantization logic to work\n",
        "first_asr_model.eval()\n",
        "audio_filepath = 'test_audio.wav'\n",
        "\n",
        "print(first_asr_model.transcribe([audio_filepath], batch_size=1)[0])"
      ],
      "metadata": {
        "id": "HY2GpUzmRlvg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "PLAY REFERENCE AUDIO"
      ],
      "metadata": {
        "id": "brmo2ZpEW_uA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sr = 16000\n",
        "import IPython.display as ipd\n",
        "# Import the data augmentation component from ASR collection\n",
        "from nemo.collections.asr.parts import perturb, segment\n",
        "def load_audio(filepath) -> segment.AudioSegment:\n",
        "    sample_segment = segment.AudioSegment.from_file(filepath, target_sr=sr)\n",
        "    return sample_segment\n",
        "\n",
        "filepath = os.path.join('/data/hxstorage/brain4x_new/data/TTS/azure_data/wavs/AudioBook/779.wav')\n",
        "sample_segment = load_audio(filepath)\n",
        "ipd.Audio(sample_segment.samples, rate=sr)"
      ],
      "metadata": {
        "id": "eAbUeJnvVkBP"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}