{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training the Spoken Language Classifier\n",
    "In this notebook, you will:\n",
    "1. Create a custom container that will be used for training and deploying\n",
    "2. Train the spoken language classifier by intiating a training job"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import libraries and load AWS credentials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -U sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "from sagemaker.estimator import Estimator\n",
    "import boto3\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "role = sagemaker.get_execution_role()\n",
    "bucket = sagemaker.Session().default_bucket()\n",
    "account_id = boto3.client('sts').get_caller_identity().get('Account')\n",
    "region = boto3.session.Session().region_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create estimator using custom container"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Build a custom container that will be used for training and deploying.  **Make sure that your role has access to ECR (attach AmazonEC2ContainerRegistryFullAccess policy to your role)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sh build_and_push.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define locations to image container, voxforge dataset in s3, and path to save models to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_uri = f'{account_id}.dkr.ecr.{region}.amazonaws.com/spoken-language-detection'\n",
    "\n",
    "s3_voxforge_prefix = os.path.join('s3://' + bucket, 'voxforge')\n",
    "s3_output_path = os.path.join('s3://' + bucket, 'models')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create estimator and define hyperparameters (default values provided below are optimal):\n",
    "* sample_rate - sample rate to load audio files\n",
    "* n_samples - number of samples to truncate audio files to (n_samples / sample_rate = duration of audio observed by model)\n",
    "* languages - list of languages to train model on\n",
    "* feature_type - type of features to extract from audio files (melspectrogram vs. MFCC)\n",
    "* normalize - whether or not to normalize the features between 0 and 1\n",
    "* standardize - whether or not to standardize the features to have zero mean and unit variance\n",
    "* standardize_mean - if standardize = True, subtract this value from the features\n",
    "* standardize_std - if standardize = True, divide the features by this value\n",
    "* random_pitch_shift - whether or not to randomly pitch shift the audio\n",
    "* pitch_shift_range - if random_pitch_shift = True, defines the range in semitones to randomly shift the audio\n",
    "* random_crop - if the audio file > n_samples, whether or not to randomly crop the audio by n_samples or take the first n_samples of the audio\n",
    "* frequency_masking - whether or not to randomly set a band of frequencies to zero (data augmentation)\n",
    "* time_masking - whether or not to randomly set a band of time steps to zero (data augmentation)\n",
    "* source_weighted_normalization - whether or not to randomly sample the training dataset such that each speaker is seen by the model equally, as well as each language.  If False, the training dataset will only be randomly sampled such that each language is seen equally\n",
    "* source_normalized_evaluation - whether or not to evaluate the model based on metrics on a per speaker basis.  This prevents the evaluation from being biased towards speakers who may have many more audio files than others\n",
    "* epochs - number of epochs to train\n",
    "* batch_size - size of each training batch\n",
    "* lr - initial learning rate\n",
    "* hidden_dim - hidden dimension before final dense layer\n",
    "* logging_iters - interval of training iterations to display training stats\n",
    "* early_stopping_epochs - training will stop if model does not improve after this many epochs\n",
    "* s3_prefix - s3 path of the VoxForge dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator = Estimator(\n",
    "    image_uri=image_uri,\n",
    "    role=role,\n",
    "    instance_count=1,\n",
    "    instance_type='ml.m5.xlarge',\n",
    "    output_path=s3_output_path,\n",
    "    max_run=5*24*60*60,\n",
    "    hyperparameters={\n",
    "        'sample_rate' : 16000,\n",
    "        'n_samples' : 80000,\n",
    "        'languages' : ['en', 'es', 'it', 'fr', 'de', 'ru'],\n",
    "        'feature_type' : 'mel',\n",
    "        'normalize' : True,\n",
    "        'standardize' : True,\n",
    "        'standardize_mean' : 0.4630,\n",
    "        'standardize_std' : 0.2031,\n",
    "        'random_pitch_shift' : True,\n",
    "        'pitch_shift_range' : (-5, 5),\n",
    "        'random_crop' : True,\n",
    "        'frequency_masking' : True,\n",
    "        'time_masking' : True,\n",
    "        'source_weighted_normalization' : True,\n",
    "        'source_normalized_evaluation' : True,\n",
    "        'epochs' : 50,\n",
    "        'batch_size' : 32,\n",
    "        'lr' : 0.0001,\n",
    "        'hidden_dim' : 512,\n",
    "        'logging_iters' : 100,\n",
    "        'early_stopping_epochs' : 5,\n",
    "        \n",
    "        's3_prefix' : s3_voxforge_prefix\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train model\n",
    "Starts the training job.  If random_pitch_shift is set to True, training can take a few days since the current python librosa implementation of pitch_shift is very slow.  However, pitch shifting will lead to the best model performance so it is not recommended to set to False.\n",
    "\n",
    "Once the training job is started, this notebook can be shutdown."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator.fit({\n",
    "    'train' : os.path.join(s3_voxforge_prefix, 'train_manifest.csv'),\n",
    "    'validation' : os.path.join(s3_voxforge_prefix, 'val_manifest.csv')\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p36",
   "language": "python",
   "name": "conda_pytorch_p36"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}