blog_post_resources/wav2vec2_inference_demo/wav2vec2_distributed_inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from jiwer import wer\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "import torchaudio\n",
    "\n",
    "import fairseq_mod\n",
    "\n",
    "from utils import Wav2VecCtc, W2lViterbiDecoder, postprocess_features, post_process_sentence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we will show you how to run distributed inference using Ray. We will use the dev-clean dataset from LibriSpeech which can be downloaded [here](https://www.openslr.org/12). We will work with the wav2vec 2.0 base model (wav2vec_small_960h) and it can be downloaded [here](https://dl.fbaipublicfiles.com/fairseq/wav2vec/wav2vec_small_960h.pt). Without further ado let's get started!\n",
    "\n",
    "### Step 1: Import and start Ray\n",
    "Note: it's better to set `num_cpus` to be much smaller than the actual number of CPUs available on the machine. To check the number of CPUs on a machine, run:\n",
    "```\n",
    "import psutil\n",
    "psutil.cpu_count(logical=True)\n",
    "```\n",
    "\n",
    "I have 48 logical CPUs on my machine and I set `num_cpus` to 20 in order to run Ray successfully"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2021-03-02 18:40:35,879\tINFO services.py:1174 -- View the Ray dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8265\u001b[39m\u001b[22m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'node_ip_address': '172.17.0.2',\n",
       " 'raylet_ip_address': '172.17.0.2',\n",
       " 'redis_address': '172.17.0.2:6379',\n",
       " 'object_store_address': '/tmp/ray/session_2021-03-02_18-40-35_397295_43186/sockets/plasma_store',\n",
       " 'raylet_socket_name': '/tmp/ray/session_2021-03-02_18-40-35_397295_43186/sockets/raylet',\n",
       " 'webui_url': '127.0.0.1:8265',\n",
       " 'session_dir': '/tmp/ray/session_2021-03-02_18-40-35_397295_43186',\n",
       " 'metrics_export_port': 62996,\n",
       " 'node_id': 'e808f4bb71ade76fec73d7d9bf97a38167cc15575bc387d922b9d276'}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ray\n",
    "ray.shutdown()\n",
    "ray.init(num_cpus=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Create the wav2vec 2.0 model, decoder and data loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = \"/home/models/wav2vec2/wav2vec_small_960h.pt\"\n",
    "data_path = \"/home/datasets\"\n",
    "target_dict = fairseq_mod.data.Dictionary.load('ltr_dict.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "w2v = torch.load(model_path)\n",
    "model = Wav2VecCtc.build_model(w2v[\"args\"], target_dict)\n",
    "model.load_state_dict(w2v[\"model\"], strict=True)\n",
    "model.eval()\n",
    "decoder = W2lViterbiDecoder(target_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev_clean_librispeech_data = torchaudio.datasets.LIBRISPEECH(data_path, url='dev-clean', download=False)\n",
    "data_loader = torch.utils.data.DataLoader(dev_clean_librispeech_data, batch_size=1, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Define a helper method which converts one audio sample into text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_data_sample(data_sample, model, decoder, target_dict):\n",
    "    encoder_input = dict()\n",
    "    feature = postprocess_features(data_sample[0][0][0], data_sample[1]).unsqueeze(0)\n",
    "    padding_mask = torch.BoolTensor(feature.size(1)).fill_(False).unsqueeze(0)\n",
    "    \n",
    "    encoder_input[\"source\"] = feature\n",
    "    encoder_input[\"padding_mask\"] = padding_mask\n",
    "    encoder_input[\"features_only\"] = True\n",
    "    encoder_input[\"mask\"] = False\n",
    "    \n",
    "    encoder_out = model(**encoder_input)\n",
    "    emissions = model.get_normalized_probs(encoder_out, log_probs=True)\n",
    "    emissions = emissions.transpose(0, 1).float().cpu().contiguous()\n",
    "    \n",
    "    decoder_out = decoder.decode(emissions)\n",
    "    hyp_pieces = target_dict.string(decoder_out[0][0][\"tokens\"].int().cpu())\n",
    "    prediction = post_process_sentence(hyp_pieces, 'letter')\n",
    "    \n",
    "    return prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Define the remote function for Ray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "@ray.remote\n",
    "def remote_process_data_sample(batch, model, generator, target_dict):\n",
    "    result = process_data_sample(batch, model, generator, target_dict)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5: Put wav2vec 2.0 model and decoder to a shared memory space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = ray.put(model)\n",
    "decoder_id = ray.put(decoder)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 6: Start the inference task for each data sample and collect predictions\n",
    "Note that Ray will distribute inference tasks in the background"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2703/2703 [00:15<00:00, 173.81it/s]\n"
     ]
    }
   ],
   "source": [
    "prediction_futures, ground_truths, start_time = [], [], time.time()\n",
    "for i, batch in enumerate(tqdm(data_loader)):\n",
    "    prediction_future = remote_process_data_sample.remote(batch, model_id, decoder_id, target_dict)\n",
    "    prediction_futures.append(prediction_future)\n",
    "    ground_truths.append(batch[2][0])\n",
    "predictions = ray.get(prediction_futures)\n",
    "inference_time = time.time() - start_time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 7: Calculate the WER score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "wer_score = wer(ground_truths, predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WER is 3.17%. Inference took 259 seconds\n"
     ]
    }
   ],
   "source": [
    "print(\"WER is {:.2f}%. Inference took {} seconds\".format(wer_score*100, int(inference_time)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}