nvidia-nemo--speech/tutorials/cloud/aws/SageMaker_ASR_Training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2e86a2b3",
   "metadata": {},
   "source": [
    "# NeMo ASR Training Using AWS SageMaker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "215e3d3c",
   "metadata": {},
   "source": [
    "In this tutorial we show how you can train a NeMo ASR Model using [Amazon SageMaker](https://docs.aws.amazon.com/sagemaker/latest/dg/whatis.html). This is meant to be a minimalistic example of how to use SageMaker with NeMo.\n",
    "\n",
    "AWS SageMaker is useful for practitioners/researchers who are familiar with training locally or on a remote instance (via SSH). SageMaker also supports multi-GPU & Multi-node.\n",
    "\n",
    "Using AWS SageMaker we train a simple Conformer CTC model using the AN4 dataset on a remote instance with a GPU (p3.2xlarge). We use S3 to store the data and our checkpoints/logs.\n",
    "\n",
    "The overall steps are:\n",
    "\n",
    "1. Setup your AWS Credentials to access SageMaker\n",
    "2. Download the source code we'll be running\n",
    "3. Setup AN4 dataset, upload data to S3\n",
    "4. Configure the training job\n",
    "5. Run training job on SageMaker\n",
    "6. Download model, (Optional) Tensorboard Logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac621da0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "5. Restart the runtime (Runtime -> Restart Runtime) for any upgraded packages to take effect\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell.\n",
    "\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 \"nemo_toolkit[all] @ git+https://github.com/NVIDIA-NeMo/Speech.git@$BRANCH\"\n",
    "\n",
    "\"\"\"\n",
    "Remember to restart the runtime for the kernel to pick up any upgraded packages (e.g. matplotlib)!\n",
    "Alternatively, you can uncomment the exit() below to crash and restart the kernel, in the case\n",
    "that you want to use the \"Run All Cells\" (or similar) option.\n",
    "\"\"\"\n",
    "# exit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61c4fbe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install sagemaker awscli"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "876f553d",
   "metadata": {},
   "source": [
    "### 1. Setup SageMaker with AWS Credentials\n",
    "\n",
    "If you haven't setup your AWS credentials, you can setup using the AWS CLI.\n",
    "You will need your access and Secret key, with permissions to use SageMaker and S3."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1328482d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws configure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01477d55",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n",
    "from pathlib import Path\n",
    "import os\n",
    "import sagemaker\n",
    "import wget\n",
    "from omegaconf import OmegaConf\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.pytorch import PyTorch\n",
    "\n",
    "from nemo.utils.notebook_utils import download_an4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "405806f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = sagemaker.Session()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d099a96",
   "metadata": {},
   "source": [
    "### 2. Download the NeMo source code\n",
    "\n",
    "SageMaker allows you to pass in your own source code, with an entrypoint script.\n",
    "\n",
    "Below we download the AWS NeMo `config.yaml` which contains our configuration, and the `speech_to_text_ctc.py` script to run training.\n",
    "\n",
    "Our folder structure will look like this:\n",
    "\n",
    "    code/\n",
    "        speech_to_text_ctc.py\n",
    "        conf/\n",
    "            config.yaml\n",
    "    data/\n",
    "        an4/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b456c57",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = Path('./an4_nemo_sagemaker/')\n",
    "code_dir = root_dir / 'code/'\n",
    "config_dir = code_dir / 'conf/'\n",
    "data_dir = root_dir / 'data/'\n",
    "\n",
    "root_dir.mkdir(exist_ok=True)\n",
    "code_dir.mkdir(exist_ok=True)\n",
    "config_dir.mkdir(exist_ok=True)\n",
    "data_dir.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "721c7b7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "config_path = str(config_dir / \"config.yaml\")\n",
    "wget.download(\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/asr/conf/conformer/conformer_ctc_char.yaml\", config_path)\n",
    "wget.download(\"https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/asr/asr_ctc/speech_to_text_ctc.py\", str(code_dir))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7934baab",
   "metadata": {},
   "source": [
    "We also create a `requirements.txt` file within our source code to install NeMo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "397d8eef",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(code_dir / 'requirements.txt', 'w') as f:\n",
    "    f.write(\"nemo_toolkit[all]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7bc7f51",
   "metadata": {},
   "source": [
    "### 2.1 Initialize SageMaker within Training Script\n",
    "\n",
    "We provide a helper function that we require to be imported and run at the top of the training script.\n",
    "\n",
    "This installs and setups DDP for you. It also alleviates having to import a custom container, and can leverage all of the SageMaker containers. Rather than running this cell, you could also manually do this in your script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e17535a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "line = \"from nemo.utils.cloud import initialize_sagemaker; initialize_sagemaker()\"\n",
    "with open(code_dir / \"speech_to_text_ctc.py\", 'r+') as f:\n",
    "    content = f.read()\n",
    "    f.seek(0, 0)\n",
    "    f.write(line.rstrip('\\r\\n') + '\\n' + content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfa2199e",
   "metadata": {},
   "source": [
    "### 3. Setup the AN4 Dataset, upload data to S3\n",
    "\n",
    "We now download our training and validation data, uploading to S3 so that SageMaker can mount our data to the instance at runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1c5a60a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# within the SageMaker container, mount_dir will be where our data is stored.\n",
    "download_an4(\n",
    "    data_dir=str(data_dir),\n",
    "    train_mount_dir=\"/opt/ml/input/data/training/\",\n",
    "    test_mount_dir=\"/opt/ml/input/data/testing/\",\n",
    ")\n",
    "\n",
    "# Upload to the default bucket\n",
    "prefix = \"an4\"\n",
    "bucket = sess.default_bucket()\n",
    "loc = sess.upload_data(path=str(data_dir), bucket=bucket, key_prefix=prefix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1de1089",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_path = \"s3://\" + sess.default_bucket() + \"/nemo-output/\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6321e3a9",
   "metadata": {},
   "source": [
    "### 4. Configure the training job\n",
    "\n",
    "Now we configure the training job, by modifying the `config.yaml` file that is stored in our source code directory.\n",
    "We pass relative directory paths for the data based on the SageMaker mount directory on the remote instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bb61640",
   "metadata": {},
   "outputs": [],
   "source": [
    "conf = OmegaConf.load(config_path)\n",
    "\n",
    "# Set Data Locations based on the mounted directory in the SageMaker instance\n",
    "conf.model.train_ds.manifest_filepath = \"/opt/ml/input/data/training/an4/train_manifest.json\"\n",
    "conf.model.validation_ds.manifest_filepath = \"/opt/ml/input/data/testing/an4/test_manifest.json\"\n",
    "conf.trainer.accelerator = \"gpu\"\n",
    "conf.trainer.max_epochs = 150\n",
    "\n",
    "# Output directory for our experiment within the SageMaker instance\n",
    "conf.exp_manager.exp_dir=\"/opt/ml/model/\"\n",
    "\n",
    "# Create a Small Variant of the Conformer Model\n",
    "conf.model.encoder.n_layers = 8\n",
    "conf.model.n_heads = 4\n",
    "conf.model.spec_augment.time_masks = 5\n",
    "\n",
    "# Set Optimizer parameters\n",
    "conf.model.optim.lr = 2.0 # by default we using Noam scheduling, the LR is a multiplier \n",
    "\n",
    "OmegaConf.save(conf, config_dir / 'config.yaml')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "959da702",
   "metadata": {},
   "source": [
    "### 5. Run training on SageMaker\n",
    "\n",
    "Pass the path of the training and validation data on S3 + the output directory on S3 to the PyTorch estimator, and call fit with the appropriate bucket locations for the training and testing data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d2e44e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "channels = {\"training\": loc, \"testing\": loc}\n",
    "\n",
    "role = get_execution_role()\n",
    "\n",
    "# Set to True to enable SageMaker to run locally\n",
    "local_mode = False\n",
    "\n",
    "if local_mode:\n",
    "    instance_type = \"local_gpu\"\n",
    "else:\n",
    "    instance_type = \"ml.p3.2xlarge\"\n",
    "\n",
    "est = PyTorch(\n",
    "    entry_point=\"speech_to_text_ctc.py\", # the script we want to run\n",
    "    source_dir=str(code_dir), # where our conf/script is\n",
    "    role=role,\n",
    "    instance_type=instance_type,\n",
    "    instance_count=1,\n",
    "    framework_version=\"1.12.0\", # version of PyTorch\n",
    "    py_version=\"py38\",\n",
    "    volume_size=250,\n",
    "    output_path=output_path,\n",
    "    hyperparameters={'config-path': 'conf'},\n",
    ")\n",
    "\n",
    "est.fit(inputs=channels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2be67b8e",
   "metadata": {},
   "source": [
    "### 6. Download model, (Optional) Tensorboard Logs\n",
    "\n",
    "SageMaker stores our models/logs within a tar file after training has finished. These can be obtained from S3 like below.\n",
    "\n",
    "We also visualize the training logs. We suggest using an external logger (such as W&B) to track training progress during the run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81cd58d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "key = est.model_data.replace(\"s3://\" + sess.default_bucket() + '/', '')\n",
    "\n",
    "sess.boto_session.client(\"s3\", region_name=sess.boto_region_name).download_file(\n",
    "    Bucket=sess.default_bucket(), Key=key, Filename='model.tar.gz',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "656d53f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!tar -xvzf model.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62018b6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%tensorboard --logdir ./ --host 0.0.0.0"
   ]
  }
 ],
 "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
}