demo/denas/asr/DENAS_ASR_DEMO.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a5a0f2a5",
   "metadata": {},
   "source": [
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/intel/e2eAIOK/blob/main/demo/denas/asr/DENAS_ASR_DEMO.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4682c66d",
   "metadata": {},
   "source": [
    "# AIOK DE-NAS ASR DEMO\n",
    "\n",
    "DE-NAS is a multi-model, hardware-aware, train-free NAS to construct compact model architectures for target platform directly. DE-NAS includes CNN-based search space for CV domain and Transformer-based search space for CV/NLP/ASR domains, and leverages hardware-aware train-free scoring method to evaluate the performance of the candidate architecture without training.\n",
    "\n",
    "This demo mainly introduces ASR integration with DE-NAS to search lighter, faster, higher performance transformer-based ASR model in a training-free way."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "337f3e2b",
   "metadata": {},
   "source": [
    "# Content\n",
    "* [Overview](#Overview)\n",
    "    * [DE-NAS on ASR Domain](#DE-NAS-on-ASR-Domain)\n",
    "    * [Performance](#Performance)\n",
    "* [Getting Started](#Getting-Started)\n",
    "    * [1. Environment Setup](#1.-Environment-Setup)\n",
    "    * [2. Workflow Prepare](#2.-Workflow-Prepare)\n",
    "    * [3. Data Prepare](#3.-Data-Prepare)\n",
    "    * [4. Search](#4.-Launch-Search)\n",
    "    * [5. Train](#5.-Train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d377abc",
   "metadata": {},
   "source": [
    "# Overview"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8c1685c8",
   "metadata": {},
   "source": [
    "## DE-NAS on ASR Domain\n",
    "\n",
    "Recently, Transformer has achieved remarkable success in several automatic speech recognition tasks. The progresses are highly relevant to the architecture design, then it is worthwhile to propose Transformer based Neural Architecture Search to search for better automatically. We will propose an unified effective method to synaptic diversity of MSA(multi-head self-attention) and synaptic saliency of MLP, which are the basic component of transformer.\n",
    "\n",
    "Transformer based search space consists of attention layer, layer normalization and feed forward layer, the search space can be controled by setting network depth, number attention heads, MLP layer ratio and layer dimension.\n",
    "\n",
    "<center>\n",
    "<img src=\"./img/asr_search_space.png\" width=\"60%\"/><figure>DE-NAS ASR Search Space and Supernet</figure>\n",
    "</center>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "421df1d6",
   "metadata": {},
   "source": [
    "## Performance\n",
    "\n",
    "<center>\n",
    "<img src=\"./img/denas_asr_perf.png\" width=\"600\"/>\n",
    "</center>\n",
    "\n",
    "* Testing methodology\n",
    "    * Dataset: LibriSpeech, Metrics: WER 5.8%\n",
    "    * Baseline: RNN-T model \n",
    "    * Early stop at WER 5.8%\n",
    "* DE-NAS ASR searched model delivered 59.12x training speedup over stock model (RNN-T).\n",
    "* Distributed training delivered 3.81x speedup with HW scaling from 1 node to 4 nodes"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "a76528ee",
   "metadata": {},
   "source": [
    "# Getting Started\n",
    "\n",
    "Notes: in order to run this demo, please follow `Environment Setup`, `Workflow Prepare` and `Data Prepare` section for pre-requirements.\n",
    "\n",
    "## 1. Environment Setup\n",
    "\n",
    "### (Option 1) Use Pip Install"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52b6b47b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "pip install e2eAIOK-denas --pre\n",
    "pip install torchsummary joblib"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "03f3d2b7",
   "metadata": {},
   "source": [
    "### (Option 2) Use Docker\n",
    "\n",
    "``` bash\n",
    "# Setup ENV\n",
    "git clone https://github.com/intel/e2eAIOK.git\n",
    "cd e2eAIOK\n",
    "python3 scripts/start_e2eaiok_docker.py -b pytorch112 -w ${host0} ${host1} ${host2} ${host3} --proxy \"\"\n",
    "# Enter Docker\n",
    "sshpass -p docker ssh ${host0} -p 12347\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6df081af",
   "metadata": {},
   "source": [
    "## 2. Workflow Prepare\n",
    "\n",
    "### search configuration\n",
    "\n",
    "```yaml\n",
    "# conf for transformer based asr\n",
    "model_type: asr\n",
    "search_engine: RandomSearchEngine #supported search engine are Random/Evolutionary/SigoptSearchEngine\n",
    "batch_size: 32\n",
    "random_max_epochs: 10 #random search max epochs\n",
    "\n",
    "#evolutionary search engine configs\n",
    "max_epochs: 10\n",
    "select_num: 50\n",
    "population_num: 50\n",
    "m_prob: 0.2\n",
    "s_prob: 0.4\n",
    "crossover_num: 25\n",
    "mutation_num: 25\n",
    "\n",
    "#searched model parameter limit\n",
    "max_param_limits: 40\n",
    "min_param_limits: 1\n",
    "\n",
    "supernet_cfg: ../../conf/denas/asr/supernet_large.conf\n",
    "img_size: 224\n",
    "seed: 0\n",
    "\n",
    "#enable/disable NAS scores\n",
    "expressivity_weight: 0\n",
    "complexity_weight: 0\n",
    "diversity_weight: 1\n",
    "saliency_weight: 1\n",
    "latency_weight: 0\n",
    "```\n",
    "\n",
    "### supernet and search space\n",
    "\n",
    "```yaml\n",
    "SUPERNET:\n",
    "  MLP_RATIO: 4.0\n",
    "  NUM_HEADS: 4\n",
    "  EMBED_DIM: 512\n",
    "  DEPTH: 12\n",
    "SEARCH_SPACE:\n",
    "  MLP_RATIO:\n",
    "    - 3.0\n",
    "    - 3.5\n",
    "    - 4.0\n",
    "    - 4.5\n",
    "    - 5.0\n",
    "  NUM_HEADS:\n",
    "    - 2\n",
    "    - 3\n",
    "    - 4\n",
    "  DEPTH:\n",
    "    - 5\n",
    "    - 6\n",
    "    - 7\n",
    "    - 8\n",
    "    - 9\n",
    "    - 10\n",
    "    - 11\n",
    "    - 12\n",
    "  EMBED_DIM:\n",
    "    - 192\n",
    "    - 216\n",
    "    - 240\n",
    "    - 324\n",
    "    - 384\n",
    "    - 444\n",
    "```\n",
    "\n",
    "### training configuration\n",
    "\n",
    "```yaml\n",
    "#edit /home/vmagent/app/e2eaiok/conf/denas/asr/e2eaiok_denas_train.conf \n",
    "train_csv: \"/home/vmagent/app/dataset/LibriSpeech/dev-clean.csv\"\n",
    "valid_csv: \"/home/vmagent/app/dataset/LibriSpeech/dev-clean.csv\"\n",
    "test_csv: \"/home/vmagent/app/dataset/LibriSpeech/dev-clean.csv\"\n",
    "tokenizer_ckpt: \"/home/vmagent/app/dataset/LibriSpeech/tokenizer.ckpt\"\n",
    "train_epochs: 1\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9f3a982",
   "metadata": {},
   "source": [
    "## 3. Data Prepare\n",
    "\n",
    "``` bash\n",
    "# Download Dataset\n",
    "# Download and unzip dataset from https://www.openslr.org/12 to /home/vmagent/app/dataset/LibriSpeech\n",
    "# Download tokenizer from https://huggingface.co/speechbrain/asr-transformer-transformerlm-librispeech/blob/main/tokenizer.ckpt to /home/vmagent/app/dataset/LibriSpeech\n",
    "\n",
    "# Process audio data\n",
    "cd ${e2eaiok_install_dir}/e2eAIOK/DeNas/asr\n",
    "conda activate pytorch\n",
    "bash scripts/preprocess_librispeech.sh\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8ac394b",
   "metadata": {},
   "source": [
    "## 4. Search\n",
    "\n",
    "Launch DENAS search on asr domain based on configs in `${e2eAIOK_install_dir}/conf/denas/asr/e2eaiok_denas_asr.conf`, searched best model structure will be saved in `/home/vmagent/app/e2eaiok/e2eAIOK/DeNas/best_model_structure.txt`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4cbf72bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "03/23/2023 08:38:34 - INFO - DENAS -   epoch = 0\n",
      "03/23/2023 08:38:35 - INFO - DENAS -   random 1/1 structure (5, 3.0, 5.0, 4.0, 5.0, 3.0, 3, 3, 3, 4, 2, 216) nas_score 0.6014334317005705 params 11.750592\n",
      "03/23/2023 08:38:35 - INFO - DENAS -   random_num = 1\n",
      "03/23/2023 08:38:36 - INFO - DENAS -   mutation 1/1 structure (5, 4.5, 4.5, 4.0, 4.5, 3.0, 3, 3, 3, 4, 2, 216) nas_score 0.29840024128498044 params 11.797356\n",
      "03/23/2023 08:38:36 - INFO - DENAS -   mutation_num = 1\n",
      "03/23/2023 08:38:36 - INFO - DENAS -   crossover_num = 0\n",
      "03/23/2023 08:38:36 - INFO - DENAS -   best structure (5, 3.0, 5.0, 4.0, 5.0, 3.0, 3, 3, 3, 4, 2, 216) nas_score 0.6014334317005705 params 11.750592\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DE-NAS completed, best structure is (5, 3.0, 5.0, 4.0, 5.0, 3.0, 3, 3, 3, 4, 2, 216)\n"
     ]
    }
   ],
   "source": [
    "import yaml\n",
    "from easydict import EasyDict as edict\n",
    "from e2eAIOK.DeNas.asr.supernet_asr import TransformerASRSuper\n",
    "from e2eAIOK.DeNas.search.SearchEngineFactory import SearchEngineFactory\n",
    "\n",
    "# create common settings\n",
    "settings = {}\n",
    "settings[\"domain\"] = \"asr\"\n",
    "# load search settings\n",
    "with open(\"/home/vmagent/app/e2eaiok/conf/denas/asr/e2eaiok_denas_asr.conf\") as f:\n",
    "    conf = yaml.load(f, Loader=yaml.FullLoader)\n",
    "settings.update(conf)\n",
    "settings[\"max_epochs\"] = 1\n",
    "settings[\"population_num\"] = 1\n",
    "settings[\"crossover_num\"] = 1\n",
    "settings[\"mutation_num\"] = 1\n",
    "params = edict(settings)\n",
    "\n",
    "# create supernet and search space\n",
    "super_net = TransformerASRSuper\n",
    "search_space = {'num_heads': params.SEARCH_SPACE.NUM_HEADS, 'mlp_ratio': params.SEARCH_SPACE.MLP_RATIO, 'embed_dim': params.SEARCH_SPACE.EMBED_DIM , 'depth': params.SEARCH_SPACE.DEPTH}\n",
    "\n",
    "# create search engine and launch search\n",
    "searcher = SearchEngineFactory.create_search_engine(params = params, super_net = super_net, search_space = search_space)\n",
    "searcher.search()\n",
    "# get best searched structure\n",
    "best_structure = searcher.get_best_structures()\n",
    "print(f\"DE-NAS completed, best structure is {best_structure}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1555538c",
   "metadata": {},
   "source": [
    "## 5. Train\n",
    "\n",
    "Load searched best model in `/home/vmagent/app/e2eaiok/e2eAIOK/DeNas/best_model_structure.txt` and launch training with training configuration in `${e2eAIOK_install_dir}/conf/denas/asr/e2eaiok_denas_train_asr.conf`\n",
    "> Note: Bellow training script is just for demonstration, and uses small sampled dataset and runs a small iterations. For actual performance result, please refer to [performance](#performance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "50a13385",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "03/23/2023 08:49:03 - INFO - Trainer -   building model\n",
      "03/23/2023 08:49:03 - INFO - Trainer -   model created\n",
      "03/23/2023 08:49:03 - INFO - Trainer -   Trainer config: {'domain': 'asr', 'seed': '74443', 'output_folder': 'results/transformer/74443', 'save_folder': 'results/transformer/74443/save', 'device': 'cpu', 'dist_backend': 'gloo', 'mode': 'train', 'best_model_structure': 'best_model_structure.txt', 'data_folder': '/home/vmagent/app/dataset/LibriSpeech', 'skip_prep': False, 'train_csv': '/home/vmagent/app/dataset/LibriSpeech/train-clean-100.csv', 'valid_csv': '/home/vmagent/app/dataset/LibriSpeech/dev-clean.csv', 'test_csv': '/home/vmagent/app/dataset/LibriSpeech/test-clean.csv', 'tokenizer_ckpt': '/home/vmagent/app/dataset/LibriSpeech/tokenizer.ckpt', 'ckpt_interval_minutes': 30, 'train_epochs': 1, 'eval_epochs': 1, 'train_batch_size': 32, 'eval_batch_size': 1, 'num_workers': 1, 'ctc_weight': 0.3, 'grad_accumulation_factor': 1, 'max_grad_norm': 5.0, 'loss_reduction': 'batchmean', 'sorting': 'random', 'metric_threshold': 25, 'lr_adam': 0.001, 'sample_rate': 16000, 'n_fft': 400, 'n_mels': 80, 'input_shape': [8, 10, 80], 'num_blocks': 3, 'num_layers_per_block': 1, 'out_channels': [64, 64, 64], 'kernel_sizes': [5, 5, 1], 'strides': [2, 2, 1], 'residuals': [False, False, True], 'input_size': 1280, 'd_model': 216, 'encoder_heads': [3, 3, 3, 4, 2], 'nhead': 4, 'num_encoder_layers': 5, 'num_decoder_layers': 6, 'mlp_ratio': [3.0, 5.0, 4.0, 5.0, 3.0], 'd_ffn': 2048, 'transformer_dropout': 0.1, 'output_neurons': 5000, 'blank_index': 0, 'label_smoothing': 0.0, 'pad_index': 0, 'bos_index': 1, 'eos_index': 2, 'min_decode_ratio': 0.0, 'max_decode_ratio': 1.0, 'valid_search_interval': 10, 'valid_beam_size': 10, 'test_beam_size': 66, 'lm_weight': 0.6, 'ctc_weight_decode': 0.4, 'n_warmup_steps': 2500, 'augmentation': {'time_warp': False, 'time_warp_window': 5, 'time_warp_mode': 'bicubic', 'freq_mask': True, 'n_freq_mask': 4, 'time_mask': True, 'n_time_mask': 4, 'replace_with_zero': False, 'freq_mask_width': 15, 'time_mask_width': 20}, 'speed_perturb': True, 'compute_features': {'sample_rate': 16000, 'n_fft': 400, 'n_mels': 80}}\n",
      "03/23/2023 08:49:07 - INFO - Trainer -   epoch: 1, step: 1|3, time: 2.95s, loss: 1089.2393798828125, avg_loss: 1089.2394, lr: 0.001\n",
      "03/23/2023 08:49:09 - INFO - Trainer -   epoch: 1, step: 2|3, time: 2.24s, loss: 918.0582275390625, avg_loss: 1003.6488, lr: 4e-07\n",
      "03/23/2023 08:49:11 - INFO - Trainer -   epoch: 1, step: 3|3, time: 2.04s, loss: 876.6215209960938, avg_loss: 961.3064, lr: 8e-07\n",
      "03/23/2023 08:49:12 - INFO - Trainer -   epoch: 1, time: 8.78s, avg_loss: 961.3064\n",
      "03/23/2023 08:49:15 - INFO - Trainer -   epoch: 1, time: 3.6744914054870605, wer: 141.82608695652175, avg_loss: 379.1408659007452\n",
      "03/23/2023 08:49:15 - INFO - Trainer -   Evaluate time:3.6786766052246094\n",
      "03/23/2023 08:49:15 - INFO - Trainer -   Epoch 1 training time:12.49078106880188\n",
      "03/23/2023 08:49:15 - INFO - Trainer -   Total time:12.491520166397095\n",
      "03/23/2023 08:49:15 - INFO - Trainer -   Trainer complete\n"
     ]
    }
   ],
   "source": [
    "import yaml\n",
    "from easydict import EasyDict as edict\n",
    "import sentencepiece as sp\n",
    "import torch\n",
    "from e2eAIOK.DeNas.asr.model_builder_denas_asr import ModelBuilderASRDeNas\n",
    "from e2eAIOK.common.trainer.data.asr.data_builder_librispeech import DataBuilderLibriSpeech\n",
    "from e2eAIOK.DeNas.asr.trainer.schedulers import NoamScheduler\n",
    "from e2eAIOK.DeNas.asr.trainer.losses import ctc_loss, kldiv_loss\n",
    "from e2eAIOK.DeNas.asr.utils.metric_stats import ErrorRateStats\n",
    "from e2eAIOK.DeNas.asr.asr_trainer import ASRTrainer\n",
    "\n",
    "# create common settings\n",
    "settings = {}\n",
    "settings[\"domain\"] = \"asr\"\n",
    "# load training settings\n",
    "with open(\"/home/vmagent/app/e2eaiok/conf/denas/asr/e2eaiok_denas_train_asr.conf\") as f:\n",
    "    conf = yaml.load(f, Loader=yaml.FullLoader)\n",
    "settings.update(conf)\n",
    "settings[\"train_epochs\"] = 1\n",
    "settings[\"best_model_structure\"] = \"best_model_structure.txt\"\n",
    "cfg = edict(settings)\n",
    "\n",
    "# create ASR model builder and create ASR model\n",
    "model = ModelBuilderASRDeNas(cfg).create_model()\n",
    "tokenizer = sp.SentencePieceProcessor()\n",
    "# get training and evaluation dataloader\n",
    "train_dataloader, eval_dataloader = DataBuilderLibriSpeech(cfg, tokenizer).get_dataloader()\n",
    "# create optimizer\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=cfg[\"lr_adam\"], betas=(0.9, 0.98), eps=0.000000001)\n",
    "criterion = {\"ctc_loss\": ctc_loss, \"seq_loss\": kldiv_loss}\n",
    "scheduler = NoamScheduler(lr_initial=cfg[\"lr_adam\"], n_warmup_steps=cfg[\"n_warmup_steps\"])\n",
    "metric = ErrorRateStats()\n",
    "# create ASR trainer\n",
    "trainer = ASRTrainer(cfg, model, train_dataloader, eval_dataloader, optimizer, criterion, scheduler, metric, tokenizer)\n",
    "# start model training and evaluation\n",
    "trainer.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21329653",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}