examples/chainer_simple/chainer_simple.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optuna example with Chainer MNIST on SageMaker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "After you create an AWS environment by the [CloudFormation template](https://github.com/aws-samples/amazon-sagemaker-optuna-hpo-blog/blob/master/template/optuna-template.yaml), install Optuna and MySQL connector to the notebook kernel, obtain parameters from the CloudFormation Outputs, and get DB secrets from AWS Secrets Manager. Please modify the `'<your_cfn_stack_name>'` to your CloudFormation stack name, which you can find at [AWS Management Console](https://us-east-1.console.aws.amazon.com/cloudformation/home?region=us-east-1#/stacks). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install optuna\n",
    "!pip install mysql-connector-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3 # AWS Python SDK\n",
    "import numpy as np\n",
    "import optuna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# obtain parameters from CloudFormation Outputs\n",
    "stack_name = '<your_cfn_stack_name>'\n",
    "\n",
    "client = boto3.client('cloudformation')\n",
    "outputs = client.describe_stacks(StackName=stack_name)['Stacks'][0]['Outputs']\n",
    "\n",
    "host = [out['OutputValue'] for out in outputs if out['OutputKey'] == 'ClusterEndpoint'][0].split(':')[0]\n",
    "db_name = [out['OutputValue'] for out in outputs if out['OutputKey'] == 'DatabaseName'][0]\n",
    "secret_name = [out['OutputValue'] for out in outputs if out['OutputKey'] == 'DBSecretArn'][0].split(':')[-1].split('-')[0]\n",
    "\n",
    "subnets = [out['OutputValue'] for out in outputs if out['OutputKey'] == 'PrivateSubnets'][0].split(',')\n",
    "security_group_ids = [out['OutputValue'] for out in outputs if out['OutputKey'] == 'SageMakerSecurityGroup'][0].split(',')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Call AWS Secrets Manager\n",
    "from src.secrets import get_secret\n",
    "region_name = boto3.session.Session().region_name\n",
    "secret = get_secret(secret_name, region_name)\n",
    "\n",
    "# MySQL-connector-python    \n",
    "db = 'mysql+mysqlconnector://{}:{}@{}/{}'.format(secret['username'], secret['password'], host, db_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup\n",
    "from sagemaker import get_execution_role\n",
    "import sagemaker\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "\n",
    "# This role retrieves the SageMaker-compatible role used by this notebook instance.\n",
    "role = get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train\n",
    "We demonstrate an Optuna example [`chainer_simple.py`](https://github.com/pfnet/optuna/blob/master/examples/chainer_simple.py) migrated to Amazon SageMaker. First, put the data to Amazon S3. Then, create a [Chainer estimator](https://sagemaker.readthedocs.io/en/stable/sagemaker.chainer.html#sagemaker.chainer.estimator.Chainer). The training will be invoked by the `fit` method (in parallel here). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create study in RDS/Aurora\n",
    "study_name = 'chainer-simple'\n",
    "optuna.study.create_study(storage=db, study_name=study_name, direction='maximize', load_if_exists=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare data\n",
    "import chainer\n",
    "import os\n",
    "import shutil\n",
    "import numpy as np\n",
    "\n",
    "N_TRAIN_EXAMPLES = 3000\n",
    "N_TEST_EXAMPLES = 1000\n",
    "\n",
    "rng = np.random.RandomState(0)\n",
    "train, test = chainer.datasets.get_mnist()\n",
    "\n",
    "train = chainer.datasets.SubDataset(\n",
    "    train, 0, N_TRAIN_EXAMPLES, order=rng.permutation(len(train)))\n",
    "test = chainer.datasets.SubDataset(\n",
    "    test, 0, N_TEST_EXAMPLES, order=rng.permutation(len(test)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = np.array([element[0] for element in train])\n",
    "train_labels = np.array([element[1] for element in train])\n",
    "\n",
    "test_data = np.array([element[0] for element in test])\n",
    "test_labels = np.array([element[1] for element in test])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# upload to Amazon S3\n",
    "try:\n",
    "    os.makedirs('/tmp/data/train_mnist')\n",
    "    os.makedirs('/tmp/data/test_mnist')\n",
    "    np.savez('/tmp/data/train_mnist/train.npz', data=train_data, labels=train_labels)\n",
    "    np.savez('/tmp/data/test_mnist/test.npz', data=test_data, labels=test_labels)\n",
    "    train_input = sagemaker_session.upload_data(\n",
    "                      path=os.path.join('/tmp', 'data', 'train_mnist'),\n",
    "                      key_prefix='notebook/chainer_mnist/train')\n",
    "    test_input = sagemaker_session.upload_data(\n",
    "                      path=os.path.join('/tmp', 'data', 'test_mnist'),\n",
    "                      key_prefix='notebook/chainer_mnist/test')\n",
    "finally:\n",
    "    shutil.rmtree('/tmp/data')\n",
    "print('training data at %s' % train_input)\n",
    "print('test data at %s' % test_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup SageMaker Chainer estimator\n",
    "from sagemaker.chainer.estimator import Chainer\n",
    "\n",
    "chainer_estimator = Chainer(entry_point='chainer_simple.py',\n",
    "                            source_dir=\"src\",\n",
    "                            framework_version='5.0.0', \n",
    "                            role=role,\n",
    "                            sagemaker_session=sagemaker_session,\n",
    "                            subnets=subnets,\n",
    "                            security_group_ids=security_group_ids,\n",
    "                            train_instance_count=1,\n",
    "                            train_instance_type='ml.c5.xlarge',\n",
    "                            hyperparameters={\n",
    "                                'host': host, \n",
    "                                'db-name': db_name, \n",
    "                                'db-secret': secret_name, \n",
    "                                'study-name': study_name, \n",
    "                                'n-trials': 25, \n",
    "                                'region-name': region_name\n",
    "                            })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# HPO in parallel\n",
    "max_parallel_jobs = 4\n",
    "\n",
    "for j in range(max_parallel_jobs-1):\n",
    "    chainer_estimator.fit({'train': train_input, 'test': test_input}, wait=False)\n",
    "chainer_estimator.fit({'train': train_input, 'test': test_input})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# obtain results\n",
    "study = optuna.study.load_study(study_name=study_name, storage=db)\n",
    "\n",
    "df = study.trials_dataframe()\n",
    "\n",
    "# optuna.visualization.plot_intermediate_values(study)\n",
    "ax = df['user_attrs']['validation/main/accuracy'].plot()\n",
    "ax.set_xlabel('Number of trials')\n",
    "ax.set_ylabel('Validation accuracy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deploy\n",
    "Create an API endopint for inference with the best model we explored in the HPO. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.chainer import ChainerModel\n",
    "\n",
    "best_model_data = os.path.join(chainer_estimator.output_path, study.best_trial.user_attrs['job_name'], 'output/model.tar.gz')\n",
    "best_model = ChainerModel(model_data=best_model_data, \n",
    "                          role=role,\n",
    "                          entry_point='chainer_simple.py', \n",
    "                          source_dir=\"src\")\n",
    "\n",
    "predictor = best_model.deploy(instance_type=\"ml.m4.xlarge\", initial_instance_count=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "num_samples = 5\n",
    "indices = random.sample(range(test_data.shape[0] - 1), num_samples)\n",
    "images, labels = test_data[indices], test_labels[indices]\n",
    "\n",
    "for i in range(num_samples):\n",
    "    plt.subplot(1,num_samples,i+1)\n",
    "    plt.imshow(images[i].reshape(28, 28), cmap='gray')\n",
    "    plt.title(labels[i])\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction = predictor.predict(images)\n",
    "predicted_label = prediction.argmax(axis=1)\n",
    "print('The predicted labels are: {}'.format(predicted_label))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cleanup\n",
    "Delete the API endpoint. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.delete_endpoint()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_chainer_p36",
   "language": "python",
   "name": "conda_chainer_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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}