sagemaker-python-sdk/mxnet_mnist/mxnet_mnist.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "executive-freight",
   "metadata": {
    "papermill": {
     "duration": 0.011048,
     "end_time": "2021-06-08T00:18:15.323922",
     "exception": false,
     "start_time": "2021-06-08T00:18:15.312874",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Training and hosting SageMaker Models using the Apache MXNet Module API\n",
    "\n",
    "The [SageMaker Python SDK](https://sagemaker.readthedocs.io/en/stable) makes it easy to train and deploy Apache MXNet models. In this example, we train a simple neural network using the Apache MXNet [Module API](https://mxnet.apache.org/api/python/module/module.html) and the MNIST dataset. The MNIST dataset is widely used for handwritten digit classification, and consists of 70,000 labeled 28x28 pixel grayscale images of hand-written digits. The dataset is split into 60,000 training images and 10,000 test images. There are 10 classes (one for each of the 10 digits). The task at hand is to train a model using the 60,000 training images and subsequently test its classification accuracy on the 10,000 test images.\n",
    "\n",
    "### Setup\n",
    "\n",
    "First we define a few variables that are needed later in this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "curious-brown",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-08T00:18:15.350043Z",
     "iopub.status.busy": "2021-06-08T00:18:15.349490Z",
     "iopub.status.idle": "2021-06-08T00:18:16.951723Z",
     "shell.execute_reply": "2021-06-08T00:18:16.951201Z"
    },
    "isConfigCell": true,
    "papermill": {
     "duration": 1.616997,
     "end_time": "2021-06-08T00:18:16.951850",
     "exception": false,
     "start_time": "2021-06-08T00:18:15.334853",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sagemaker import get_execution_role\n",
    "from sagemaker.session import Session\n",
    "\n",
    "# S3 bucket for saving code and model artifacts.\n",
    "# Feel free to specify a different bucket here if you wish.\n",
    "bucket = Session().default_bucket()\n",
    "\n",
    "# Bucket location where your custom code will be saved in the tar.gz format.\n",
    "custom_code_upload_location = \"s3://{}/mxnet-mnist-example/code\".format(bucket)\n",
    "\n",
    "# Bucket location where results of model training are saved.\n",
    "model_artifacts_location = \"s3://{}/mxnet-mnist-example/artifacts\".format(bucket)\n",
    "\n",
    "# IAM execution role that gives SageMaker access to resources in your AWS account.\n",
    "# We can use the SageMaker Python SDK to get the role from our notebook environment.\n",
    "role = get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "emerging-mounting",
   "metadata": {
    "papermill": {
     "duration": 0.011079,
     "end_time": "2021-06-08T00:18:16.974197",
     "exception": false,
     "start_time": "2021-06-08T00:18:16.963118",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### The training script\n",
    "\n",
    "The `mnist.py` script provides all the code we need for training and hosting a SageMaker model. The script also checkpoints the model at the end of every epoch and saves the model graph, params and optimizer state in the folder `/opt/ml/checkpoints`. If the folder path does not exist then it skips checkpointing. The script we use is adaptated from Apache MXNet [MNIST tutorial](https://mxnet.incubator.apache.org/tutorials/python/mnist.html).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "atomic-objective",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-08T00:18:17.000620Z",
     "iopub.status.busy": "2021-06-08T00:18:17.000089Z",
     "iopub.status.idle": "2021-06-08T00:18:17.822741Z",
     "shell.execute_reply": "2021-06-08T00:18:17.822251Z"
    },
    "papermill": {
     "duration": 0.837588,
     "end_time": "2021-06-08T00:18:17.822860",
     "exception": false,
     "start_time": "2021-06-08T00:18:16.985272",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36margparse\u001b[39;49;00m\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mgzip\u001b[39;49;00m\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mjson\u001b[39;49;00m\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mlogging\u001b[39;49;00m\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mos\u001b[39;49;00m\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mstruct\u001b[39;49;00m\n",
      "\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mmxnet\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mmx\u001b[39;49;00m\n",
      "\u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mnumpy\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mnp\u001b[39;49;00m\n",
      "\n",
      "\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mload_data\u001b[39;49;00m(path):\n",
      "    \u001b[34mwith\u001b[39;49;00m gzip.open(find_file(path, \u001b[33m\"\u001b[39;49;00m\u001b[33mlabels.gz\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)) \u001b[34mas\u001b[39;49;00m flbl:\n",
      "        struct.unpack(\u001b[33m\"\u001b[39;49;00m\u001b[33m>II\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m, flbl.read(\u001b[34m8\u001b[39;49;00m))\n",
      "        labels = np.fromstring(flbl.read(), dtype=np.int8)\n",
      "    \u001b[34mwith\u001b[39;49;00m gzip.open(find_file(path, \u001b[33m\"\u001b[39;49;00m\u001b[33mimages.gz\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)) \u001b[34mas\u001b[39;49;00m fimg:\n",
      "        _, _, rows, cols = struct.unpack(\u001b[33m\"\u001b[39;49;00m\u001b[33m>IIII\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m, fimg.read(\u001b[34m16\u001b[39;49;00m))\n",
      "        images = np.fromstring(fimg.read(), dtype=np.uint8).reshape(\u001b[36mlen\u001b[39;49;00m(labels), rows, cols)\n",
      "        images = images.reshape(images.shape[\u001b[34m0\u001b[39;49;00m], \u001b[34m1\u001b[39;49;00m, \u001b[34m28\u001b[39;49;00m, \u001b[34m28\u001b[39;49;00m).astype(np.float32) / \u001b[34m255\u001b[39;49;00m\n",
      "    \u001b[34mreturn\u001b[39;49;00m labels, images\n",
      "\n",
      "\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mfind_file\u001b[39;49;00m(root_path, file_name):\n",
      "    \u001b[34mfor\u001b[39;49;00m root, dirs, files \u001b[35min\u001b[39;49;00m os.walk(root_path):\n",
      "        \u001b[34mif\u001b[39;49;00m file_name \u001b[35min\u001b[39;49;00m files:\n",
      "            \u001b[34mreturn\u001b[39;49;00m os.path.join(root, file_name)\n",
      "\n",
      "\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mbuild_graph\u001b[39;49;00m():\n",
      "    data = mx.sym.var(\u001b[33m\"\u001b[39;49;00m\u001b[33mdata\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\n",
      "    data = mx.sym.flatten(data=data)\n",
      "    fc1 = mx.sym.FullyConnected(data=data, num_hidden=\u001b[34m128\u001b[39;49;00m)\n",
      "    act1 = mx.sym.Activation(data=fc1, act_type=\u001b[33m\"\u001b[39;49;00m\u001b[33mrelu\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\n",
      "    fc2 = mx.sym.FullyConnected(data=act1, num_hidden=\u001b[34m64\u001b[39;49;00m)\n",
      "    act2 = mx.sym.Activation(data=fc2, act_type=\u001b[33m\"\u001b[39;49;00m\u001b[33mrelu\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\n",
      "    fc3 = mx.sym.FullyConnected(data=act2, num_hidden=\u001b[34m10\u001b[39;49;00m)\n",
      "    \u001b[34mreturn\u001b[39;49;00m mx.sym.SoftmaxOutput(data=fc3, name=\u001b[33m\"\u001b[39;49;00m\u001b[33msoftmax\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\n",
      "\n",
      "\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mget_training_context\u001b[39;49;00m(num_gpus):\n",
      "    \u001b[34mif\u001b[39;49;00m num_gpus:\n",
      "        \u001b[34mreturn\u001b[39;49;00m [mx.gpu(i) \u001b[34mfor\u001b[39;49;00m i \u001b[35min\u001b[39;49;00m \u001b[36mrange\u001b[39;49;00m(num_gpus)]\n",
      "    \u001b[34melse\u001b[39;49;00m:\n",
      "        \u001b[34mreturn\u001b[39;49;00m mx.cpu()\n",
      "\n",
      "\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mtrain\u001b[39;49;00m(\n",
      "    batch_size,\n",
      "    epochs,\n",
      "    learning_rate,\n",
      "    num_gpus,\n",
      "    training_channel,\n",
      "    testing_channel,\n",
      "    hosts,\n",
      "    current_host,\n",
      "    model_dir,\n",
      "):\n",
      "    checkpoints_dir = \u001b[33m\"\u001b[39;49;00m\u001b[33m/opt/ml/checkpoints\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m\n",
      "    checkpoints_enabled = os.path.exists(checkpoints_dir)\n",
      "\n",
      "    (train_labels, train_images) = load_data(training_channel)\n",
      "    (test_labels, test_images) = load_data(testing_channel)\n",
      "    \u001b[37m# Data parallel training - shard the data so each host\u001b[39;49;00m\n",
      "    \u001b[37m# only trains on a subset of the total data.\u001b[39;49;00m\n",
      "    shard_size = \u001b[36mlen\u001b[39;49;00m(train_images) // \u001b[36mlen\u001b[39;49;00m(hosts)\n",
      "    \u001b[34mfor\u001b[39;49;00m i, host \u001b[35min\u001b[39;49;00m \u001b[36menumerate\u001b[39;49;00m(hosts):\n",
      "        \u001b[34mif\u001b[39;49;00m host == current_host:\n",
      "            start = shard_size * i\n",
      "            end = start + shard_size\n",
      "            \u001b[34mbreak\u001b[39;49;00m\n",
      "\n",
      "    train_iter = mx.io.NDArrayIter(\n",
      "        train_images[start:end], train_labels[start:end], batch_size, shuffle=\u001b[34mTrue\u001b[39;49;00m\n",
      "    )\n",
      "    val_iter = mx.io.NDArrayIter(test_images, test_labels, batch_size)\n",
      "\n",
      "    logging.getLogger().setLevel(logging.DEBUG)\n",
      "\n",
      "    kvstore = \u001b[33m\"\u001b[39;49;00m\u001b[33mlocal\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m \u001b[34mif\u001b[39;49;00m \u001b[36mlen\u001b[39;49;00m(hosts) == \u001b[34m1\u001b[39;49;00m \u001b[34melse\u001b[39;49;00m \u001b[33m\"\u001b[39;49;00m\u001b[33mdist_sync\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m\n",
      "\n",
      "    mlp_model = mx.mod.Module(symbol=build_graph(), context=get_training_context(num_gpus))\n",
      "\n",
      "    checkpoint_callback = \u001b[34mNone\u001b[39;49;00m\n",
      "    \u001b[34mif\u001b[39;49;00m checkpoints_enabled:\n",
      "        \u001b[37m# Create a checkpoint callback that checkpoints the model params and\u001b[39;49;00m\n",
      "        \u001b[37m# the optimizer state at the given path after every epoch.\u001b[39;49;00m\n",
      "        checkpoint_callback = mx.callback.module_checkpoint(\n",
      "            mlp_model, os.path.join(checkpoints_dir, \u001b[33m\"\u001b[39;49;00m\u001b[33mmnist\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m), period=\u001b[34m1\u001b[39;49;00m, save_optimizer_states=\u001b[34mTrue\u001b[39;49;00m\n",
      "        )\n",
      "    mlp_model.fit(\n",
      "        train_iter,\n",
      "        eval_data=val_iter,\n",
      "        kvstore=kvstore,\n",
      "        optimizer=\u001b[33m\"\u001b[39;49;00m\u001b[33msgd\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m,\n",
      "        optimizer_params={\u001b[33m\"\u001b[39;49;00m\u001b[33mlearning_rate\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m: learning_rate},\n",
      "        eval_metric=\u001b[33m\"\u001b[39;49;00m\u001b[33macc\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m,\n",
      "        epoch_end_callback=checkpoint_callback,\n",
      "        batch_end_callback=mx.callback.Speedometer(batch_size, \u001b[34m100\u001b[39;49;00m),\n",
      "        num_epoch=epochs,\n",
      "    )\n",
      "\n",
      "    \u001b[34mif\u001b[39;49;00m current_host == hosts[\u001b[34m0\u001b[39;49;00m]:\n",
      "        save(model_dir, mlp_model)\n",
      "\n",
      "\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32msave\u001b[39;49;00m(model_dir, model):\n",
      "    model.symbol.save(os.path.join(model_dir, \u001b[33m\"\u001b[39;49;00m\u001b[33mmodel-symbol.json\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m))\n",
      "    model.save_params(os.path.join(model_dir, \u001b[33m\"\u001b[39;49;00m\u001b[33mmodel-0000.params\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m))\n",
      "\n",
      "    signature = [\n",
      "        {\u001b[33m\"\u001b[39;49;00m\u001b[33mname\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m: data_desc.name, \u001b[33m\"\u001b[39;49;00m\u001b[33mshape\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m: [dim \u001b[34mfor\u001b[39;49;00m dim \u001b[35min\u001b[39;49;00m data_desc.shape]}\n",
      "        \u001b[34mfor\u001b[39;49;00m data_desc \u001b[35min\u001b[39;49;00m model.data_shapes\n",
      "    ]\n",
      "    \u001b[34mwith\u001b[39;49;00m \u001b[36mopen\u001b[39;49;00m(os.path.join(model_dir, \u001b[33m\"\u001b[39;49;00m\u001b[33mmodel-shapes.json\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m), \u001b[33m\"\u001b[39;49;00m\u001b[33mw\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m) \u001b[34mas\u001b[39;49;00m f:\n",
      "        json.dump(signature, f)\n",
      "\n",
      "\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mparse_args\u001b[39;49;00m():\n",
      "    parser = argparse.ArgumentParser()\n",
      "\n",
      "    parser.add_argument(\u001b[33m\"\u001b[39;49;00m\u001b[33m--batch-size\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m100\u001b[39;49;00m)\n",
      "    parser.add_argument(\u001b[33m\"\u001b[39;49;00m\u001b[33m--epochs\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mint\u001b[39;49;00m, default=\u001b[34m10\u001b[39;49;00m)\n",
      "    parser.add_argument(\u001b[33m\"\u001b[39;49;00m\u001b[33m--learning-rate\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mfloat\u001b[39;49;00m, default=\u001b[34m0.1\u001b[39;49;00m)\n",
      "\n",
      "    parser.add_argument(\u001b[33m\"\u001b[39;49;00m\u001b[33m--model-dir\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, default=os.environ[\u001b[33m\"\u001b[39;49;00m\u001b[33mSM_MODEL_DIR\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m])\n",
      "    parser.add_argument(\u001b[33m\"\u001b[39;49;00m\u001b[33m--train\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, default=os.environ[\u001b[33m\"\u001b[39;49;00m\u001b[33mSM_CHANNEL_TRAIN\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m])\n",
      "    parser.add_argument(\u001b[33m\"\u001b[39;49;00m\u001b[33m--test\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, default=os.environ[\u001b[33m\"\u001b[39;49;00m\u001b[33mSM_CHANNEL_TEST\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m])\n",
      "\n",
      "    parser.add_argument(\u001b[33m\"\u001b[39;49;00m\u001b[33m--current-host\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mstr\u001b[39;49;00m, default=os.environ[\u001b[33m\"\u001b[39;49;00m\u001b[33mSM_CURRENT_HOST\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m])\n",
      "    parser.add_argument(\u001b[33m\"\u001b[39;49;00m\u001b[33m--hosts\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m, \u001b[36mtype\u001b[39;49;00m=\u001b[36mlist\u001b[39;49;00m, default=json.loads(os.environ[\u001b[33m\"\u001b[39;49;00m\u001b[33mSM_HOSTS\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m]))\n",
      "\n",
      "    \u001b[34mreturn\u001b[39;49;00m parser.parse_args()\n",
      "\n",
      "\n",
      "\u001b[37m### NOTE: this function cannot use MXNet\u001b[39;49;00m\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mneo_preprocess\u001b[39;49;00m(payload, content_type):\n",
      "    \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mio\u001b[39;49;00m\n",
      "    \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mlogging\u001b[39;49;00m\n",
      "\n",
      "    \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mnumpy\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mnp\u001b[39;49;00m\n",
      "\n",
      "    logging.info(\u001b[33m\"\u001b[39;49;00m\u001b[33mInvoking user-defined pre-processing function\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\n",
      "\n",
      "    \u001b[34mif\u001b[39;49;00m content_type != \u001b[33m\"\u001b[39;49;00m\u001b[33mapplication/vnd+python.numpy+binary\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m:\n",
      "        \u001b[34mraise\u001b[39;49;00m \u001b[36mRuntimeError\u001b[39;49;00m(\u001b[33m\"\u001b[39;49;00m\u001b[33mContent type must be application/vnd+python.numpy+binary\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\n",
      "\n",
      "    f = io.BytesIO(payload)\n",
      "    \u001b[34mreturn\u001b[39;49;00m np.load(f)\n",
      "\n",
      "\n",
      "\u001b[37m### NOTE: this function cannot use MXNet\u001b[39;49;00m\n",
      "\u001b[34mdef\u001b[39;49;00m \u001b[32mneo_postprocess\u001b[39;49;00m(result):\n",
      "    \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mjson\u001b[39;49;00m\n",
      "    \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mlogging\u001b[39;49;00m\n",
      "\n",
      "    \u001b[34mimport\u001b[39;49;00m \u001b[04m\u001b[36mnumpy\u001b[39;49;00m \u001b[34mas\u001b[39;49;00m \u001b[04m\u001b[36mnp\u001b[39;49;00m\n",
      "\n",
      "    logging.info(\u001b[33m\"\u001b[39;49;00m\u001b[33mInvoking user-defined post-processing function\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m)\n",
      "\n",
      "    \u001b[37m# Softmax (assumes batch size 1)\u001b[39;49;00m\n",
      "    result = np.squeeze(result)\n",
      "    result_exp = np.exp(result - np.max(result))\n",
      "    result = result_exp / np.sum(result_exp)\n",
      "\n",
      "    response_body = json.dumps(result.tolist())\n",
      "    content_type = \u001b[33m\"\u001b[39;49;00m\u001b[33mapplication/json\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m\n",
      "\n",
      "    \u001b[34mreturn\u001b[39;49;00m response_body, content_type\n",
      "\n",
      "\n",
      "\u001b[34mif\u001b[39;49;00m \u001b[31m__name__\u001b[39;49;00m == \u001b[33m\"\u001b[39;49;00m\u001b[33m__main__\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m:\n",
      "    args = parse_args()\n",
      "    num_gpus = \u001b[36mint\u001b[39;49;00m(os.environ[\u001b[33m\"\u001b[39;49;00m\u001b[33mSM_NUM_GPUS\u001b[39;49;00m\u001b[33m\"\u001b[39;49;00m])\n",
      "\n",
      "    train(\n",
      "        args.batch_size,\n",
      "        args.epochs,\n",
      "        args.learning_rate,\n",
      "        num_gpus,\n",
      "        args.train,\n",
      "        args.test,\n",
      "        args.hosts,\n",
      "        args.current_host,\n",
      "        args.model_dir,\n",
      "    )\n"
     ]
    }
   ],
   "source": [
    "!pygmentize mnist.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "chicken-share",
   "metadata": {
    "papermill": {
     "duration": 0.011903,
     "end_time": "2021-06-08T00:18:17.847179",
     "exception": false,
     "start_time": "2021-06-08T00:18:17.835276",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### SageMaker's MXNet estimator class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "occasional-fountain",
   "metadata": {
    "papermill": {
     "duration": 0.011774,
     "end_time": "2021-06-08T00:18:17.870755",
     "exception": false,
     "start_time": "2021-06-08T00:18:17.858981",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "The SageMaker ```MXNet``` estimator allows us to run single machine or distributed training in SageMaker, using CPU or GPU-based instances.\n",
    "\n",
    "When we create the estimator, we pass in the filename of our training script, the name of our IAM execution role, and the S3 locations we defined in the setup section. We also provide a few other parameters. ``train_instance_count`` and ``train_instance_type`` determine the number and type of SageMaker instances that will be used for the training job. The ``hyperparameters`` parameter is a ``dict`` of values that will be passed to your training script -- you can see how to access these values in the ``mnist.py`` script above.\n",
    "\n",
    "For this example, we will choose one ``ml.m4.xlarge`` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "choice-consolidation",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-08T00:18:17.899339Z",
     "iopub.status.busy": "2021-06-08T00:18:17.898827Z",
     "iopub.status.idle": "2021-06-08T00:18:17.968039Z",
     "shell.execute_reply": "2021-06-08T00:18:17.968486Z"
    },
    "papermill": {
     "duration": 0.085933,
     "end_time": "2021-06-08T00:18:17.968634",
     "exception": false,
     "start_time": "2021-06-08T00:18:17.882701",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "distributions has been renamed in sagemaker>=2.\n",
      "See: https://sagemaker.readthedocs.io/en/stable/v2.html for details.\n"
     ]
    }
   ],
   "source": [
    "from sagemaker.mxnet import MXNet\n",
    "\n",
    "mnist_estimator = MXNet(\n",
    "    entry_point=\"mnist.py\",\n",
    "    role=role,\n",
    "    output_path=model_artifacts_location,\n",
    "    code_location=custom_code_upload_location,\n",
    "    instance_count=1,\n",
    "    instance_type=\"ml.m4.xlarge\",\n",
    "    framework_version=\"1.4.1\",\n",
    "    py_version=\"py3\",\n",
    "    #distribution={\"parameter_server\": {\"enabled\": True}},\n",
    "    hyperparameters={\"learning-rate\": 0.1},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tropical-anniversary",
   "metadata": {
    "papermill": {
     "duration": 0.012999,
     "end_time": "2021-06-08T00:18:17.994933",
     "exception": false,
     "start_time": "2021-06-08T00:18:17.981934",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Running the Training Job"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "generic-alarm",
   "metadata": {
    "papermill": {
     "duration": 0.013019,
     "end_time": "2021-06-08T00:18:18.021081",
     "exception": false,
     "start_time": "2021-06-08T00:18:18.008062",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "After we've constructed our MXNet object, we can fit it using data stored in S3. Below we run SageMaker training on two input channels: **train** and **test**.\n",
    "\n",
    "During training, SageMaker makes this data stored in S3 available in the local filesystem where the mnist script is running. The ```mnist.py``` script simply loads the train and test data from disk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "documentary-passport",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-08T00:18:18.061478Z",
     "iopub.status.busy": "2021-06-08T00:18:18.059374Z",
     "iopub.status.idle": "2021-06-08T00:23:00.740466Z",
     "shell.execute_reply": "2021-06-08T00:23:00.739990Z"
    },
    "papermill": {
     "duration": 282.706559,
     "end_time": "2021-06-08T00:23:00.740588",
     "exception": false,
     "start_time": "2021-06-08T00:18:18.034029",
     "status": "completed"
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2021-06-09 21:55:39 Starting - Starting the training job...\n",
      "2021-06-09 21:55:41 Starting - Launching requested ML instances......\n",
      "2021-06-09 21:56:50 Starting - Preparing the instances for training......\n",
      "2021-06-09 21:57:59 Downloading - Downloading input data...\n",
      "2021-06-09 21:58:24 Training - Downloading the training image..\u001b[34m2021-06-09 21:58:47,583 sagemaker-containers INFO     Imported framework sagemaker_mxnet_container.training\u001b[0m\n",
      "\u001b[34m2021-06-09 21:58:47,587 sagemaker-containers INFO     No GPUs detected (normal if no gpus installed)\u001b[0m\n",
      "\u001b[34m2021-06-09 21:58:47,603 sagemaker_mxnet_container.training INFO     MXNet training environment: {'SM_HOSTS': '[\"algo-1\"]', 'SM_NETWORK_INTERFACE_NAME': 'eth0', 'SM_HPS': '{\"learning-rate\":0.1}', 'SM_USER_ENTRY_POINT': 'mnist.py', 'SM_FRAMEWORK_PARAMS': '{\"sagemaker_parameter_server_enabled\":true}', 'SM_RESOURCE_CONFIG': '{\"current_host\":\"algo-1\",\"hosts\":[\"algo-1\"],\"network_interface_name\":\"eth0\"}', 'SM_INPUT_DATA_CONFIG': '{\"test\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"},\"train\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"}}', 'SM_OUTPUT_DATA_DIR': '/opt/ml/output/data', 'SM_CHANNELS': '[\"test\",\"train\"]', 'SM_CURRENT_HOST': 'algo-1', 'SM_MODULE_NAME': 'mnist', 'SM_LOG_LEVEL': '20', 'SM_FRAMEWORK_MODULE': 'sagemaker_mxnet_container.training:main', 'SM_INPUT_DIR': '/opt/ml/input', 'SM_INPUT_CONFIG_DIR': '/opt/ml/input/config', 'SM_OUTPUT_DIR': '/opt/ml/output', 'SM_NUM_CPUS': '4', 'SM_NUM_GPUS': '0', 'SM_MODEL_DIR': '/opt/ml/model', 'SM_MODULE_DIR': 's3://sagemaker-us-west-2-688520471316/mxnet-mnist-example/code/mxnet-training-2021-06-09-21-55-39-493/source/sourcedir.tar.gz', 'SM_TRAINING_ENV': '{\"additional_framework_parameters\":{\"sagemaker_parameter_server_enabled\":true},\"channel_input_dirs\":{\"test\":\"/opt/ml/input/data/test\",\"train\":\"/opt/ml/input/data/train\"},\"current_host\":\"algo-1\",\"framework_module\":\"sagemaker_mxnet_container.training:main\",\"hosts\":[\"algo-1\"],\"hyperparameters\":{\"learning-rate\":0.1},\"input_config_dir\":\"/opt/ml/input/config\",\"input_data_config\":{\"test\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"},\"train\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"}},\"input_dir\":\"/opt/ml/input\",\"is_master\":true,\"job_name\":\"mxnet-training-2021-06-09-21-55-39-493\",\"log_level\":20,\"master_hostname\":\"algo-1\",\"model_dir\":\"/opt/ml/model\",\"module_dir\":\"s3://sagemaker-us-west-2-688520471316/mxnet-mnist-example/code/mxnet-training-2021-06-09-21-55-39-493/source/sourcedir.tar.gz\",\"module_name\":\"mnist\",\"network_interface_name\":\"eth0\",\"num_cpus\":4,\"num_gpus\":0,\"output_data_dir\":\"/opt/ml/output/data\",\"output_dir\":\"/opt/ml/output\",\"output_intermediate_dir\":\"/opt/ml/output/intermediate\",\"resource_config\":{\"current_host\":\"algo-1\",\"hosts\":[\"algo-1\"],\"network_interface_name\":\"eth0\"},\"user_entry_point\":\"mnist.py\"}', 'SM_USER_ARGS': '[\"--learning-rate\",\"0.1\"]', 'SM_OUTPUT_INTERMEDIATE_DIR': '/opt/ml/output/intermediate', 'SM_CHANNEL_TEST': '/opt/ml/input/data/test', 'SM_CHANNEL_TRAIN': '/opt/ml/input/data/train', 'SM_HP_LEARNING-RATE': '0.1'}\u001b[0m\n",
      "\u001b[34m2021-06-09 21:58:49,033 sagemaker_mxnet_container.training INFO     Starting distributed training task\u001b[0m\n",
      "\n",
      "2021-06-09 21:58:46 Training - Training image download completed. Training in progress.\u001b[34m2021-06-09 21:59:50,941 sagemaker-containers INFO     Module mnist does not provide a setup.py. \u001b[0m\n",
      "\u001b[34mGenerating setup.py\u001b[0m\n",
      "\u001b[34m2021-06-09 21:59:50,942 sagemaker-containers INFO     Generating setup.cfg\u001b[0m\n",
      "\u001b[34m2021-06-09 21:59:50,942 sagemaker-containers INFO     Generating MANIFEST.in\u001b[0m\n",
      "\u001b[34m2021-06-09 21:59:50,942 sagemaker-containers INFO     Installing module with the following command:\u001b[0m\n",
      "\u001b[34m/usr/local/bin/python3.6 -m pip install -U . \u001b[0m\n",
      "\u001b[34mProcessing /opt/ml/code\u001b[0m\n",
      "\u001b[34mInstalling collected packages: mnist\n",
      "  Running setup.py install for mnist: started\u001b[0m\n",
      "\u001b[34m    Running setup.py install for mnist: finished with status 'done'\u001b[0m\n",
      "\u001b[34mSuccessfully installed mnist-1.0.0\u001b[0m\n",
      "\u001b[34mWARNING: You are using pip version 19.1.1, however version 21.1.2 is available.\u001b[0m\n",
      "\u001b[34mYou should consider upgrading via the 'pip install --upgrade pip' command.\u001b[0m\n",
      "\u001b[34m2021-06-09 21:59:52,869 sagemaker-containers INFO     No GPUs detected (normal if no gpus installed)\u001b[0m\n",
      "\u001b[34m2021-06-09 21:59:52,887 sagemaker-containers INFO     Invoking user script\n",
      "\u001b[0m\n",
      "\u001b[34mTraining Env:\n",
      "\u001b[0m\n",
      "\u001b[34m{\n",
      "    \"additional_framework_parameters\": {\n",
      "        \"sagemaker_parameter_server_enabled\": true\n",
      "    },\n",
      "    \"channel_input_dirs\": {\n",
      "        \"test\": \"/opt/ml/input/data/test\",\n",
      "        \"train\": \"/opt/ml/input/data/train\"\n",
      "    },\n",
      "    \"current_host\": \"algo-1\",\n",
      "    \"framework_module\": \"sagemaker_mxnet_container.training:main\",\n",
      "    \"hosts\": [\n",
      "        \"algo-1\"\n",
      "    ],\n",
      "    \"hyperparameters\": {\n",
      "        \"learning-rate\": 0.1\n",
      "    },\n",
      "    \"input_config_dir\": \"/opt/ml/input/config\",\n",
      "    \"input_data_config\": {\n",
      "        \"test\": {\n",
      "            \"TrainingInputMode\": \"File\",\n",
      "            \"S3DistributionType\": \"FullyReplicated\",\n",
      "            \"RecordWrapperType\": \"None\"\n",
      "        },\n",
      "        \"train\": {\n",
      "            \"TrainingInputMode\": \"File\",\n",
      "            \"S3DistributionType\": \"FullyReplicated\",\n",
      "            \"RecordWrapperType\": \"None\"\n",
      "        }\n",
      "    },\n",
      "    \"input_dir\": \"/opt/ml/input\",\n",
      "    \"is_master\": true,\n",
      "    \"job_name\": \"mxnet-training-2021-06-09-21-55-39-493\",\n",
      "    \"log_level\": 20,\n",
      "    \"master_hostname\": \"algo-1\",\n",
      "    \"model_dir\": \"/opt/ml/model\",\n",
      "    \"module_dir\": \"s3://sagemaker-us-west-2-688520471316/mxnet-mnist-example/code/mxnet-training-2021-06-09-21-55-39-493/source/sourcedir.tar.gz\",\n",
      "    \"module_name\": \"mnist\",\n",
      "    \"network_interface_name\": \"eth0\",\n",
      "    \"num_cpus\": 4,\n",
      "    \"num_gpus\": 0,\n",
      "    \"output_data_dir\": \"/opt/ml/output/data\",\n",
      "    \"output_dir\": \"/opt/ml/output\",\n",
      "    \"output_intermediate_dir\": \"/opt/ml/output/intermediate\",\n",
      "    \"resource_config\": {\n",
      "        \"current_host\": \"algo-1\",\n",
      "        \"hosts\": [\n",
      "            \"algo-1\"\n",
      "        ],\n",
      "        \"network_interface_name\": \"eth0\"\n",
      "    },\n",
      "    \"user_entry_point\": \"mnist.py\"\u001b[0m\n",
      "\u001b[34m}\n",
      "\u001b[0m\n",
      "\u001b[34mEnvironment variables:\n",
      "\u001b[0m\n",
      "\u001b[34mSM_HOSTS=[\"algo-1\"]\u001b[0m\n",
      "\u001b[34mSM_NETWORK_INTERFACE_NAME=eth0\u001b[0m\n",
      "\u001b[34mSM_HPS={\"learning-rate\":0.1}\u001b[0m\n",
      "\u001b[34mSM_USER_ENTRY_POINT=mnist.py\u001b[0m\n",
      "\u001b[34mSM_FRAMEWORK_PARAMS={\"sagemaker_parameter_server_enabled\":true}\u001b[0m\n",
      "\u001b[34mSM_RESOURCE_CONFIG={\"current_host\":\"algo-1\",\"hosts\":[\"algo-1\"],\"network_interface_name\":\"eth0\"}\u001b[0m\n",
      "\u001b[34mSM_INPUT_DATA_CONFIG={\"test\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"},\"train\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"}}\u001b[0m\n",
      "\u001b[34mSM_OUTPUT_DATA_DIR=/opt/ml/output/data\u001b[0m\n",
      "\u001b[34mSM_CHANNELS=[\"test\",\"train\"]\u001b[0m\n",
      "\u001b[34mSM_CURRENT_HOST=algo-1\u001b[0m\n",
      "\u001b[34mSM_MODULE_NAME=mnist\u001b[0m\n",
      "\u001b[34mSM_LOG_LEVEL=20\u001b[0m\n",
      "\u001b[34mSM_FRAMEWORK_MODULE=sagemaker_mxnet_container.training:main\u001b[0m\n",
      "\u001b[34mSM_INPUT_DIR=/opt/ml/input\u001b[0m\n",
      "\u001b[34mSM_INPUT_CONFIG_DIR=/opt/ml/input/config\u001b[0m\n",
      "\u001b[34mSM_OUTPUT_DIR=/opt/ml/output\u001b[0m\n",
      "\u001b[34mSM_NUM_CPUS=4\u001b[0m\n",
      "\u001b[34mSM_NUM_GPUS=0\u001b[0m\n",
      "\u001b[34mSM_MODEL_DIR=/opt/ml/model\u001b[0m\n",
      "\u001b[34mSM_MODULE_DIR=s3://sagemaker-us-west-2-688520471316/mxnet-mnist-example/code/mxnet-training-2021-06-09-21-55-39-493/source/sourcedir.tar.gz\u001b[0m\n",
      "\u001b[34mSM_TRAINING_ENV={\"additional_framework_parameters\":{\"sagemaker_parameter_server_enabled\":true},\"channel_input_dirs\":{\"test\":\"/opt/ml/input/data/test\",\"train\":\"/opt/ml/input/data/train\"},\"current_host\":\"algo-1\",\"framework_module\":\"sagemaker_mxnet_container.training:main\",\"hosts\":[\"algo-1\"],\"hyperparameters\":{\"learning-rate\":0.1},\"input_config_dir\":\"/opt/ml/input/config\",\"input_data_config\":{\"test\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"},\"train\":{\"RecordWrapperType\":\"None\",\"S3DistributionType\":\"FullyReplicated\",\"TrainingInputMode\":\"File\"}},\"input_dir\":\"/opt/ml/input\",\"is_master\":true,\"job_name\":\"mxnet-training-2021-06-09-21-55-39-493\",\"log_level\":20,\"master_hostname\":\"algo-1\",\"model_dir\":\"/opt/ml/model\",\"module_dir\":\"s3://sagemaker-us-west-2-688520471316/mxnet-mnist-example/code/mxnet-training-2021-06-09-21-55-39-493/source/sourcedir.tar.gz\",\"module_name\":\"mnist\",\"network_interface_name\":\"eth0\",\"num_cpus\":4,\"num_gpus\":0,\"output_data_dir\":\"/opt/ml/output/data\",\"output_dir\":\"/opt/ml/output\",\"output_intermediate_dir\":\"/opt/ml/output/intermediate\",\"resource_config\":{\"current_host\":\"algo-1\",\"hosts\":[\"algo-1\"],\"network_interface_name\":\"eth0\"},\"user_entry_point\":\"mnist.py\"}\u001b[0m\n",
      "\u001b[34mSM_USER_ARGS=[\"--learning-rate\",\"0.1\"]\u001b[0m\n",
      "\u001b[34mSM_OUTPUT_INTERMEDIATE_DIR=/opt/ml/output/intermediate\u001b[0m\n",
      "\u001b[34mSM_CHANNEL_TEST=/opt/ml/input/data/test\u001b[0m\n",
      "\u001b[34mSM_CHANNEL_TRAIN=/opt/ml/input/data/train\u001b[0m\n",
      "\u001b[34mSM_HP_LEARNING-RATE=0.1\u001b[0m\n",
      "\u001b[34mPYTHONPATH=/usr/local/bin:/usr/local/lib/python36.zip:/usr/local/lib/python3.6:/usr/local/lib/python3.6/lib-dynload:/usr/local/lib/python3.6/site-packages\n",
      "\u001b[0m\n",
      "\u001b[34mInvoking script with the following command:\n",
      "\u001b[0m\n",
      "\u001b[34m/usr/local/bin/python3.6 -m mnist --learning-rate 0.1\n",
      "\n",
      "\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[0] Batch [0-100]#011Speed: 47287.58 samples/sec#011accuracy=0.110990\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[0] Batch [100-200]#011Speed: 54047.38 samples/sec#011accuracy=0.110000\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[0] Batch [200-300]#011Speed: 55089.63 samples/sec#011accuracy=0.110500\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[0] Batch [300-400]#011Speed: 51203.88 samples/sec#011accuracy=0.111900\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[0] Batch [400-500]#011Speed: 53620.52 samples/sec#011accuracy=0.114500\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[0] Train-accuracy=0.131200\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[0] Time cost=1.200\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[0] Validation-accuracy=0.363200\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[1] Batch [0-100]#011Speed: 46955.54 samples/sec#011accuracy=0.474455\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[1] Batch [100-200]#011Speed: 48841.06 samples/sec#011accuracy=0.662100\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[1] Batch [200-300]#011Speed: 54751.28 samples/sec#011accuracy=0.766200\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[1] Batch [300-400]#011Speed: 51805.90 samples/sec#011accuracy=0.800000\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[1] Batch [400-500]#011Speed: 53854.95 samples/sec#011accuracy=0.828500\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[1] Train-accuracy=0.728033\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[1] Time cost=1.141\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[1] Validation-accuracy=0.841800\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[2] Batch [0-100]#011Speed: 33069.71 samples/sec#011accuracy=0.853960\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[2] Batch [100-200]#011Speed: 39941.72 samples/sec#011accuracy=0.876500\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[2] Batch [200-300]#011Speed: 39199.43 samples/sec#011accuracy=0.885800\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[2] Batch [300-400]#011Speed: 51825.04 samples/sec#011accuracy=0.896700\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[2] Batch [400-500]#011Speed: 54412.47 samples/sec#011accuracy=0.904300\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[2] Train-accuracy=0.887317\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[2] Time cost=1.366\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[2] Validation-accuracy=0.910200\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[3] Batch [0-100]#011Speed: 44400.05 samples/sec#011accuracy=0.920594\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[3] Batch [100-200]#011Speed: 53323.57 samples/sec#011accuracy=0.926800\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[3] Batch [200-300]#011Speed: 58114.34 samples/sec#011accuracy=0.930900\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[3] Batch [300-400]#011Speed: 58136.58 samples/sec#011accuracy=0.928700\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[3] Batch [400-500]#011Speed: 53841.67 samples/sec#011accuracy=0.933100\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[3] Train-accuracy=0.929700\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[3] Time cost=1.247\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[3] Validation-accuracy=0.938600\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[4] Batch [0-100]#011Speed: 48411.49 samples/sec#011accuracy=0.943168\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[4] Batch [100-200]#011Speed: 53567.72 samples/sec#011accuracy=0.944800\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[4] Batch [200-300]#011Speed: 50246.77 samples/sec#011accuracy=0.941500\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[4] Batch [300-400]#011Speed: 44958.83 samples/sec#011accuracy=0.947900\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[4] Batch [400-500]#011Speed: 41663.27 samples/sec#011accuracy=0.948200\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[4] Train-accuracy=0.946000\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[4] Time cost=1.477\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[4] Validation-accuracy=0.943700\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[5] Batch [0-100]#011Speed: 42967.82 samples/sec#011accuracy=0.951980\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[5] Batch [100-200]#011Speed: 55124.16 samples/sec#011accuracy=0.958500\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[5] Batch [200-300]#011Speed: 41605.54 samples/sec#011accuracy=0.957700\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[5] Batch [300-400]#011Speed: 39872.05 samples/sec#011accuracy=0.958000\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[5] Batch [400-500]#011Speed: 40017.75 samples/sec#011accuracy=0.957300\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[5] Train-accuracy=0.957117\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[5] Time cost=1.405\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[5] Validation-accuracy=0.959900\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[6] Batch [0-100]#011Speed: 26208.98 samples/sec#011accuracy=0.963168\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[6] Batch [100-200]#011Speed: 37885.71 samples/sec#011accuracy=0.960500\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[6] Batch [200-300]#011Speed: 45689.64 samples/sec#011accuracy=0.963600\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[6] Batch [300-400]#011Speed: 56166.17 samples/sec#011accuracy=0.967000\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[6] Batch [400-500]#011Speed: 51141.88 samples/sec#011accuracy=0.966300\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[6] Train-accuracy=0.964417\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[6] Time cost=1.438\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[6] Validation-accuracy=0.961900\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[7] Batch [0-100]#011Speed: 47852.92 samples/sec#011accuracy=0.970693\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[7] Batch [100-200]#011Speed: 53753.14 samples/sec#011accuracy=0.966600\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[7] Batch [200-300]#011Speed: 52497.57 samples/sec#011accuracy=0.968000\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[7] Batch [300-400]#011Speed: 42965.93 samples/sec#011accuracy=0.966300\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[7] Batch [400-500]#011Speed: 53491.00 samples/sec#011accuracy=0.970200\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[7] Train-accuracy=0.968967\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[7] Time cost=1.200\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[7] Validation-accuracy=0.964300\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[8] Batch [0-100]#011Speed: 39102.81 samples/sec#011accuracy=0.975347\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[8] Batch [100-200]#011Speed: 49170.00 samples/sec#011accuracy=0.969900\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[8] Batch [200-300]#011Speed: 58197.32 samples/sec#011accuracy=0.973700\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[8] Batch [300-400]#011Speed: 49250.02 samples/sec#011accuracy=0.970600\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[8] Batch [400-500]#011Speed: 47072.32 samples/sec#011accuracy=0.972700\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[8] Train-accuracy=0.972783\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[8] Time cost=1.240\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[8] Validation-accuracy=0.968300\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[9] Batch [0-100]#011Speed: 44996.50 samples/sec#011accuracy=0.978812\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[9] Batch [100-200]#011Speed: 53172.89 samples/sec#011accuracy=0.972600\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[9] Batch [200-300]#011Speed: 58175.85 samples/sec#011accuracy=0.976500\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[9] Batch [300-400]#011Speed: 58447.10 samples/sec#011accuracy=0.975300\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[9] Batch [400-500]#011Speed: 58023.90 samples/sec#011accuracy=0.976400\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[9] Train-accuracy=0.975617\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[9] Time cost=1.139\u001b[0m\n",
      "\u001b[34mINFO:root:Epoch[9] Validation-accuracy=0.965500\u001b[0m\n",
      "\u001b[34m2021-06-09 22:00:12,939 sagemaker-containers INFO     Reporting training SUCCESS\u001b[0m\n",
      "\n",
      "2021-06-09 22:00:21 Uploading - Uploading generated training model\n",
      "2021-06-09 22:00:21 Completed - Training job completed\n",
      "Training seconds: 142\n",
      "Billable seconds: 142\n",
      "CPU times: user 599 ms, sys: 40.2 ms, total: 640 ms\n",
      "Wall time: 5min 14s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "import boto3\n",
    "\n",
    "region = boto3.Session().region_name\n",
    "train_data_location = \"s3://sagemaker-sample-data-{}/mxnet/mnist/train\".format(region)\n",
    "test_data_location = \"s3://sagemaker-sample-data-{}/mxnet/mnist/test\".format(region)\n",
    "\n",
    "mnist_estimator.fit({\"train\": train_data_location, \"test\": test_data_location})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "documentary-feedback",
   "metadata": {
    "papermill": {
     "duration": 0.020269,
     "end_time": "2021-06-08T00:23:00.781496",
     "exception": false,
     "start_time": "2021-06-08T00:23:00.761227",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "### Opimtize your model with Neo API\n",
    "Neo API allows to optimize our model for a specific hardware type. When calling compile_model() function, we specify the target instance family (C5) as well as the S3 bucket to which the compiled model would be stored.\n",
    "\n",
    "#### Important. If the following command result in a permission error, scroll up and locate the value of execution role returned by get_execution_role(). The role must have access to the S3 bucket specified in output_path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "canadian-forth",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-06-08T00:23:00.828221Z",
     "iopub.status.busy": "2021-06-08T00:23:00.827688Z",
     "iopub.status.idle": "2021-06-08T00:23:00.964047Z",
     "shell.execute_reply": "2021-06-08T00:23:00.962987Z"
    },
    "papermill": {
     "duration": 0.162609,
     "end_time": "2021-06-08T00:23:00.964285",
     "exception": true,
     "start_time": "2021-06-08T00:23:00.801676",
     "status": "failed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "?..................................................!"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Defaulting to the only supported framework/algorithm version: 1.7. Ignoring framework/algorithm version: 1.8.0.\n"
     ]
    }
   ],
   "source": [
    "output_path = \"/\".join(mnist_estimator.output_path.split(\"/\")[:-1])\n",
    "neo_optimize = True\n",
    "compiled_model = mnist_estimator.compile_model(\n",
    "    target_instance_family=\"ml_m4\",\n",
    "    input_shape={\"data\": [1, 784], \"softmax_label\": [1]},\n",
    "    role=role,\n",
    "    output_path=output_path,\n",
    "    framework=\"mxnet\",\n",
    "    framework_version=\"1.8.0\",\n",
    "    \n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rotary-enforcement",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "source": [
    "### Creating an inference Endpoint\n",
    "\n",
    "After training, we use the ``MXNet estimator`` object to build and deploy an ``MXNetPredictor``. This creates a Sagemaker **Endpoint** -- a hosted prediction service that we can use to perform inference. \n",
    "\n",
    "The arguments to the ``deploy`` function allow us to set the number and type of instances that will be used for the Endpoint. These do not need to be the same as the values we used for the training job. For example, you can train a model on a set of GPU-based instances, and then deploy the Endpoint to a fleet of CPU-based instances. Here we will deploy the model to a single ``ml.m4.xlarge`` instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "closed-china",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "------!"
     ]
    }
   ],
   "source": [
    "import io\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "def numpy_bytes_serializer(data):\n",
    "    f = io.BytesIO()\n",
    "    np.save(f, data)\n",
    "    f.seek(0)\n",
    "    return f.read()\n",
    "\n",
    "serializer = None\n",
    "if neo_optimize is True:\n",
    "    serializer = numpy_bytes_serializer\n",
    "    \n",
    "predictor = compiled_model.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type='ml.m4.xlarge',\n",
    "    serializer=serializer\n",
    "    \n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "spare-visibility",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "source": [
    "# (Optional) Delete the Endpoint\n",
    "\n",
    "After you have finished with this example, remember to delete the prediction endpoint to release the instance(s) associated with it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "killing-animation",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Endpoint name: mxnet-training-ml-m4-2021-06-09-23-15-20-517\n"
     ]
    }
   ],
   "source": [
    "print(\"Endpoint name: \" + predictor.endpoint_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "pleasant-rainbow",
   "metadata": {
    "papermill": {
     "duration": null,
     "end_time": null,
     "exception": null,
     "start_time": null,
     "status": "pending"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "predictor.delete_endpoint()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Environment (conda_mxnet_p36)",
   "language": "python",
   "name": "conda_mxnet_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"
  },
  "notice": "Copyright 2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.  Licensed under the Apache License, Version 2.0 (the \"License\"). You may not use this file except in compliance with the License. A copy of the License is located at http://aws.amazon.com/apache2.0/ or in the \"license\" file accompanying this file. This file is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.",
  "papermill": {
   "default_parameters": {},
   "duration": 286.751707,
   "end_time": "2021-06-08T00:23:01.390469",
   "environment_variables": {},
   "exception": true,
   "input_path": "mxnet_mnist.ipynb",
   "output_path": "/opt/ml/processing/output/mxnet_mnist-2021-06-08-00-14-15.ipynb",
   "parameters": {
    "kms_key": "arn:aws:kms:us-west-2:521695447989:key/6e9984db-50cf-4c7e-926c-877ec47a8b25"
   },
   "start_time": "2021-06-08T00:18:14.638762",
   "version": "2.3.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}