Add disclaimer and local gpu

Author: Dan Choi

advanced_functionality/pytorch_extending_our_containers/pytorch_extending_our_containers.ipynb

@@ -46,14 +46,13 @@
     "\n",
     "Even if there is direct SDK support for your environment or framework, you may want to add additional functionality or configure your container environment differently while utilizing our container to use on SageMaker.\n",
     "\n",
-    "All of our deep learning framework containers are open source on GitHub. Based on your use case it you can easily modify the container code and build a modified version of our containers.\n",
+    "**Some of the reasons to extend a SageMaker deep learning framework container are:**\n",
+    "1. Install additional dependencies, while utilizing the same training and hosting solution.\n",
+    "2. Configure your environment.\n",
     "\n",
-    "* [SageMaker TensorFlow Container](https://github.com/aws/sagemaker-tensorflow-container)\n",
-    "* [SageMaker MXNet Container](https://github.com/aws/sagemaker-mxnet-container)\n",
-    "* [SageMaker PyTorch Container](https://github.com/aws/sagemaker-pytorch-container) \n",
-    "* [SageMaker Chainer Container](https://github.com/aws/sagemaker-chainer-container)\n",
+    "**Although it is possible to extend any of our framework containers as a parent image, the example this notebook covers is currently only intended to work with our PyTorch (0.4.0+) and Chainer (4.1.0+) containers.**\n",
     "\n",
-    "This walkthrough shows that it is quite straightforward to extend one of our containers to build your own custom container.\n",
+    "This walkthrough shows that it is quite straightforward to extend one of our containers to build your own custom container for PyTorch or Chainer.\n",
     "\n",
     "## Permissions\n",
     "\n",
@@ -84,7 +83,7 @@
     "\n",
     "If you're familiar with Docker already, you can skip ahead to the next section.\n",
     "\n",
-    "For many data scientists, Docker containers are a new technology. But they are not difficult and can significantly simply the deployment of your software packages. \n",
+    "For many data scientists, Docker containers are a new technology. But they are not difficult and can significantly simplify the deployment of your software packages. \n",
     "\n",
     "Docker provides a simple way to package arbitrary code into an _image_ that is totally self-contained. Once you have an image, you can use Docker to run a _container_ based on that image. Running a container is just like running a program on the machine except that the container creates a fully self-contained environment for the program to run. Containers are isolated from each other and from the host environment, so the way your program is set up is the way it runs, no matter where you run it.\n",
     "\n",
@@ -250,6 +249,7 @@
       "\r\n",
       "# For more information on creating a Dockerfile\r\n",
       "# https://docs.docker.com/compose/gettingstarted/#step-2-create-a-dockerfile\r\n",
+      "# https://github.com/awslabs/amazon-sagemaker-examples/master/advanced_functionality/pytorch_extending_our_containers/pytorch_extending_our_containers.ipynb\r\n",
       "# SageMaker PyTorch image\r\n",
       "FROM 520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-pytorch:0.4.0-cpu-py3\r\n",
       "\r\n",
@@ -375,7 +375,7 @@
     "To represent our training, we use the Estimator class, which needs to be configured in five steps. \n",
     "1. IAM role - our AWS execution role\n",
     "2. train_instance_count - number of instances to use for training.\n",
-    "3. train_instance_type - type of instance to use for training. For training locally, we specify `local`.\n",
+    "3. train_instance_type - type of instance to use for training. For training locally, we specify `local` or `local_gpu`.\n",
     "4. image_name - our custom PyTorch Docker image we created.\n",
     "5. hyperparameters - hyperparameters we want to pass.\n",
     "\n",
@@ -420,6 +420,24 @@
     "!/bin/bash ./utils/setup.sh"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import subprocess\n",
+    "\n",
+    "instance_type = 'local'\n",
+    "\n",
+    "if subprocess.call('nvidia-smi') == 0:\n",
+    "    ## Set type to GPU if one is present\n",
+    "    instance_type = 'local_gpu'\n",
+    "    \n",
+    "print(\"Instance type = \" + instance_type)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -430,8 +448,6 @@
     "\n",
     "hyperparameters = {'epochs': 1}\n",
     "\n",
-    "instance_type = 'local'\n",
-    "\n",
     "estimator = Estimator(role=role,\n",
     "                      train_instance_count=1,\n",
     "                      train_instance_type=instance_type,\n",