vit-pytorch-sagemaker.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image Classification using Vision Transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import cv2\n",
    "import PIL.Image as Image\n",
    "import seaborn as sns\n",
    "from pylab import rcParams\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import rc\n",
    "from matplotlib.ticker import MaxNLocator\n",
    "from glob import glob\n",
    "import shutil\n",
    "\n",
    "import torch, torchvision\n",
    "from torch import nn, optim\n",
    "import torchvision.transforms as T\n",
    "from torchvision.datasets import ImageFolder\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import sagemaker\n",
    "\n",
    "\n",
    "RANDOM_SEED = 42\n",
    "np.random.seed(RANDOM_SEED)\n",
    "torch.manual_seed(RANDOM_SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_session = sagemaker.Session()\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "prefix = \"sagemaker/pytorch-vit\"\n",
    "role = sagemaker.get_execution_role()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "The dataset used  is `Intel Image Classification`. \n",
    "It contains around 25k images of size 150x150 distributed under 6 categories.\n",
    "```\n",
    "{'buildings' -> 0,\n",
    "'forest' -> 1,\n",
    "'glacier' -> 2,\n",
    "'mountain' -> 3,\n",
    "'sea' -> 4,\n",
    "'street' -> 5 }\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from zipfile import ZipFile\n",
    "with ZipFile('data1.zip', 'r') as zipObj:\n",
    "   \n",
    "   zipObj.extractall('data1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set = './data1/seg_train/seg_train'\n",
    "test_set = './data1/seg_test/seg_test'\n",
    "pred_set = './data1/seg_pred/seg_pred'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class_names = ['buildings', 'forest', 'glacier', 'mountain', 'sea', 'street']\n",
    "class_indices = [0,1,2,3,4,5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_folders = sorted(glob(train_set + '/*'))\n",
    "len(train_folders)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_image(img_path, resize=True):\n",
    "    img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    if resize:\n",
    "        img = cv2.resize(img, (64,64), interpolation = cv2.INTER_AREA)\n",
    "    \n",
    "    return img\n",
    "\n",
    "def show_image(img_path):\n",
    "    img = load_image(img_path)\n",
    "    plt.imshow(img)\n",
    "    plt.axis('off')\n",
    "    \n",
    "def show_sign_grid(image_paths):\n",
    "    images = [load_image(img) for img in image_paths]\n",
    "    images = torch.as_tensor(images)\n",
    "    images = images.permute(0,3,1,2)\n",
    "    grid_img = torchvision.utils.make_grid(images, nrow=11)\n",
    "    plt.figure(figsize=(24,12))\n",
    "    plt.imshow(grid_img.permute(1,2,0))\n",
    "    plt.axis('off')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_images = [np.random.choice(glob(f'{tf}/*jpg')) for tf in train_folders]\n",
    "show_sign_grid(sample_images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf data\n",
    "\n",
    "DATA_DIR = Path('data')\n",
    "\n",
    "DATASETS = ['train', 'val']\n",
    "\n",
    "for ds in DATASETS:\n",
    "    for cls in class_names:\n",
    "        (DATA_DIR / ds / cls).mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i, cls_index in enumerate(class_indices):\n",
    "    image_paths = np.array(glob(f'{train_folders[cls_index]}/*jpg'))\n",
    "    class_name = class_names[i]\n",
    "    print(f'{class_name}: {len(image_paths)}')\n",
    "    np.random.shuffle(image_paths)\n",
    "    \n",
    "    ds_split = np.split(\n",
    "        image_paths,\n",
    "        indices_or_sections=[int(.8*len(image_paths)), int(.9*len(image_paths))]\n",
    "    )\n",
    "    \n",
    "    dataset_data = zip(DATASETS, ds_split)\n",
    "    for ds, images in dataset_data:\n",
    "        for img_path in images:\n",
    "            shutil.copy(img_path, f'{DATA_DIR}/{ds}/{class_name}/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_nums = [0.485, 0.456, 0.406]\n",
    "std_nums = [0.229, 0.224, 0.225]\n",
    "\n",
    "transforms = {'train': T.Compose([\n",
    "    T.RandomResizedCrop(size=224),\n",
    "    T.RandomRotation(degrees=15),\n",
    "    T.RandomHorizontalFlip(),\n",
    "    T.ToTensor(),\n",
    "    T.Normalize(mean_nums, std_nums)\n",
    "]), 'val': T.Compose([\n",
    "    T.Resize(size=224),\n",
    "    T.CenterCrop(size=224),\n",
    "    T.ToTensor(),\n",
    "    T.Normalize(mean_nums, std_nums)\n",
    "]),}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_datasets = {\n",
    "    d: ImageFolder(f'{DATA_DIR}/{d}', transforms[d]) for d in DATASETS\n",
    "}\n",
    "\n",
    "data_loaders = {\n",
    "    d: DataLoader(image_datasets[d], batch_size=16, shuffle=True, num_workers=4) for d in DATASETS\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_sizes = {d: len(image_datasets[d]) for d in DATASETS}\n",
    "class_names = image_datasets['train'].classes\n",
    "dataset_sizes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def imshow(inp, title=None):\n",
    "    inp = inp.numpy().transpose((1,2,0))\n",
    "    mean = np.array([mean_nums])\n",
    "    std = np.array([std_nums])\n",
    "    inp = std * inp + mean\n",
    "    inp = np.clip(inp,0,1)\n",
    "    plt.imshow(inp)\n",
    "    if title is not None:\n",
    "        plt.title(title)\n",
    "    plt.axis('off')\n",
    "    \n",
    "inputs, classes = next(iter(data_loaders['train']))\n",
    "out = torchvision.utils.make_grid(inputs)\n",
    "\n",
    "imshow(out, title=[class_names[x] for x in classes])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# upload data to S3\n",
    "input_path = sagemaker_session.upload_data(path='data', bucket=bucket, key_prefix=prefix)\n",
    "print('input specification (in this case, just an S3 path): {}'.format(input_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training job on Sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "from sagemaker.pytorch import PyTorch\n",
    "\n",
    "now = datetime.now()\n",
    "timestr = now.strftime(\"%m-%d-%Y-%H-%M-%S\")\n",
    "vt_training_job_name = \"vt-training-{}\".format(timestr)\n",
    "print(vt_training_job_name)\n",
    "\n",
    "estimator = PyTorch(\n",
    "    entry_point=\"vit-job.py\",\n",
    "    source_dir=\"code\",\n",
    "    role=role,\n",
    "    framework_version=\"1.6.0\",\n",
    "    py_version=\"py3\",\n",
    "    instance_count=1,  \n",
    "    instance_type=\"ml.p3.16xlarge\", \n",
    "    use_spot_instances=False,\n",
    "    debugger_hook_config=False,\n",
    "    hyperparameters={\n",
    "        \"epochs\": 5,\n",
    "        \"num_classes\": 6,\n",
    "        \"batch-size\": 256,\n",
    "    },\n",
    "    metric_definitions=[\n",
    "                   {'Name': 'validation:loss', 'Regex': 'Valid_loss = ([0-9\\\\.]+);'},\n",
    "                   {'Name': 'validation:accuracy', 'Regex': 'Valid_accuracy = ([0-9\\\\.]+);'},\n",
    "                   {'Name': 'train:accuracy', 'Regex': 'Train_accuracy = ([0-9\\\\.]+);'},\n",
    "                   {'Name': 'train:loss', 'Regex': 'Train_loss = ([0-9\\\\.]+);'},\n",
    "                ]\n",
    ")\n",
    "estimator.fit({\"training\": input_path}, wait=True, job_name=vt_training_job_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vt_training_job_name = estimator.latest_training_job.name\n",
    "print(\"Vision Transformer training job name: \", vt_training_job_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Deploy to sagemaker endpoint\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker import get_execution_role\n",
    "ENDPOINT_NAME='pytorch-inference-{}'.format(timestr)\n",
    "predictor = estimator.deploy(initial_instance_count=1, instance_type='ml.p3.2xlarge', endpoint_name=ENDPOINT_NAME)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import requests\n",
    "from IPython.display import Image \n",
    "import json\n",
    "import boto3\n",
    "import numpy as np\n",
    "\n",
    "runtime= boto3.client('runtime.sagemaker')\n",
    "client = boto3.client('sagemaker')\n",
    "\n",
    "endpoint_desc = client.describe_endpoint(EndpointName=ENDPOINT_NAME)\n",
    "print(endpoint_desc)\n",
    "print('---'*60)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predictions from Sagemaker endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "payload =  '[{\"url\":\"https://multimedia-commons.s3-us-west-2.amazonaws.com/data/images/139/019/1390196df443f2cf614f2255ae75fcf8.jpg\"},\\\n",
    "{\"url\":\"https://multimedia-commons.s3-us-west-2.amazonaws.com/data/images/139/015/1390157d4caaf290962de5c5fb4c42.jpg\"},\\\n",
    "{\"url\":\"https://multimedia-commons.s3-us-west-2.amazonaws.com/data/images/139/020/1390207be327f4c4df1259c7266473.jpg\"},\\\n",
    "{\"url\":\"https://multimedia-commons.s3-us-west-2.amazonaws.com/data/images/139/021/139021f9aed9896831bf88f349fcec.jpg\"},\\\n",
    "{\"url\":\"https://multimedia-commons.s3-us-west-2.amazonaws.com/data/images/139/028/139028d865bafa3de66568eeb499f4a6.jpg\"},\\\n",
    "{\"url\":\"https://multimedia-commons.s3-us-west-2.amazonaws.com/data/images/139/030/13903090f3c8c7a708ca69c8d5d68b2.jpg\"},\\\n",
    "{\"url\":\"https://multimedia-commons.s3-us-west-2.amazonaws.com/data/images/002/010/00201099c5bf0d794c9a951b74390.jpg\"},\\\n",
    "{\"url\":\"https://multimedia-commons.s3-us-west-2.amazonaws.com/data/images/139/136/139136bb43e41df8949f873fb44af.jpg\"},\\\n",
    "{\"url\":\"https://multimedia-commons.s3-us-west-2.amazonaws.com/data/images/139/145/1391457e4a2e25557cbf956aaee4345.jpg\"}]'\n",
    "\n",
    "payload = json.loads(payload)\n",
    "for item in payload:\n",
    "    item = json.dumps(item)\n",
    "    response = runtime.invoke_endpoint(EndpointName=ENDPOINT_NAME, \n",
    "                                   ContentType='application/json', \n",
    "                                   Body=item)\n",
    "    result = response['Body'].read()\n",
    "    result = json.loads(result)\n",
    "    print('predicted:', result[0]['prediction'])\n",
    "\n",
    "    from PIL import Image\n",
    "    import requests\n",
    "\n",
    "    input_data = json.loads(item)\n",
    "    url = input_data['url']\n",
    "    im = Image.open(requests.get(url, stream=True).raw)\n",
    "    newsize = (250, 250) \n",
    "    im1 = im.resize(newsize) \n",
    "\n",
    "    from IPython.display import Image\n",
    "    display(im1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.delete_endpoint()"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.g4dn.xlarge",
  "kernelspec": {
   "display_name": "Python 3 (PyTorch 1.6 Python 3.6 GPU Optimized)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-west-2:236514542706:image/pytorch-1.6-gpu-py36-cu110-ubuntu18.04-v3"
  },
  "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
}