content: add PyTorch Image Classification Single GPU using Vertex Training with Custom Container and Prediction with Custom TorchServe Container (#38)

2026-09-26 14:42:04 +00:00 · 2021-08-23 15:05:40 -07:00
parent af764b0d98
commit b994d5bb97
15 changed files with 1670 additions and 0 deletions
@@ -0,0 +1,31 @@
+# Copyright 2021 Google LLC. 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.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License 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.
+"""Converts image to Base64."""
+
+import base64
+
+INPUT_FILE = 'sample.jpg'
+OUTPUT_FILE = 'sample_b64.json'
+
+def convert_to_base64(image_file):
+  """Open image and convert it to base64"""
+  with open(image_file, 'rb') as input_file:
+    jpeg_bytes = base64.b64encode(input_file.read()).decode('utf-8')
+    predict_request = '{"instances" : [{"data": {"b64": "%s"}}]}' % jpeg_bytes
+    # Write JSON to file
+    with open(OUTPUT_FILE, 'w') as output_file:
+      output_file.write(predict_request)
+    return predict_request
+
+convert_to_base64(INPUT_FILE)
@@ -0,0 +1,22 @@
+FROM pytorch/torchserve:0.3.0-cpu
+
+COPY . /home/model-server/
+WORKDIR /home/model-server/
+
+USER model-server
+RUN torch-model-archiver \
+  --model-name=antandbee \
+  --version=1.0 \
+  --model-file=./model.py \
+  --serialized-file=./model/antandbee.pth \
+  --handler=./handler.py \
+  --extra-files=./index_to_name.json \
+  --export-path=./model-store \
+  -f
+
+CMD ["torchserve", \
+     "--model-store ./model-store", \
+     "--ts-config ./config.properties", \
+     "--models antandbee=antandbee.mar"]
+
+
@@ -0,0 +1,7 @@
+inference_address=http://0.0.0.0:8080
+management_address=http://0.0.0.0:8081
+metrics_address=http://0.0.0.0:8082
+number_of_netty_threads=32
+job_queue_size=1000
+model_store=/home/model-server/model-store
+service_envelope=json
@@ -0,0 +1,18 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the \"License\");
+# you may not use this file except in compliance with the License.\n",
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License 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.
+
+from ts.torch_handler.image_classifier import ImageClassifier as ImgClsfr
+
+class ImageClassifier(ImgClsfr):
+  topk = 2
@@ -0,0 +1 @@
+{"0": ["id01", "ant"], "1": ["id02", "bee"]}
@@ -0,0 +1,19 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the \"License\");
+# you may not use this file except in compliance with the License.\n",
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License 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.
+
+from torchvision.models.resnet import ResNet, BasicBlock
+
+class ImageClassifier(ResNet):
+  def __init__(self):
+    super(ImageClassifier, self).__init__(BasicBlock, [2,2,2,2], num_classes = 2)
@@ -0,0 +1,4 @@
+google-cloud-aiplatform
+torch-model-archiver
+torchserve
+captum
@@ -0,0 +1,16 @@
+FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-runtime
+
+RUN apt-get update && \
+    apt-get install -y curl gnupg && \
+    echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
+    curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg  add - && \
+    apt-get update -y && \
+    apt-get install google-cloud-sdk -y
+
+COPY . /trainer
+
+WORKDIR /trainer
+
+RUN pip install -r requirements.txt
+
+ENTRYPOINT ["python", "-m", "task"]
@@ -0,0 +1,3 @@
+torch==1.8.1
+torchvision==0.9.1
+tensorboard==2.5.0
@@ -0,0 +1,291 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the \"License\");
+# you may not use this file except in compliance with the License.\n",
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License 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.
+
+import argparse
+import copy
+import os
+import pathlib
+import time
+
+import torch
+from torch import optim
+from torch.utils.tensorboard import SummaryWriter
+
+from torchvision import datasets, models, transforms
+
+import utils
+
+def parse_args():
+
+  parser = argparse.ArgumentParser()
+
+  # Using environment variables for Cloud Storage directories
+  # see more details in https://cloud.google.com/vertex-ai/docs/training/code-requirements
+  parser.add_argument(
+      '--model-dir', default=os.getenv('AIP_MODEL_DIR'), type=str,
+      help='a Cloud Storage URI of a directory intended for saving model artifacts')
+  parser.add_argument(
+      '--tensorboard-log-dir', default=os.getenv('AIP_TENSORBOARD_LOG_DIR'), type=str,
+      help='a Cloud Storage URI of a directory intended for saving TensorBoard')
+
+  parser.add_argument(
+      '--epochs', default=25, type=int,
+      help='number of training epochs')
+  parser.add_argument(
+      '--learning-rate', default=0.001, type=float,
+      help='learning rate')
+  parser.add_argument(
+      '--momentum', default=0.9, type=float,
+      help='momentum')
+  parser.add_argument(
+      '--batch-size', default=4, type=int,
+      help='mini-batch size')
+  parser.add_argument(
+      '--num-workers', default=4, type=int,
+      help='number of workers')
+
+  parser.add_argument(
+      '--local-mode', action='store_true', help='use local mode when running on your local machine')
+
+  args = parser.parse_args()
+
+  return args
+
+def download_data(data_dir):
+
+  dataset_url = 'https://download.pytorch.org/tutorial/hymenoptera_data.zip'
+  data_dir = os.path.abspath(data_dir)
+  datasets.utils.download_url(
+      url=dataset_url,
+      root=data_dir
+  )
+
+  from_path=os.path.join(data_dir, 'hymenoptera_data.zip')
+  datasets.utils.extract_archive(
+      from_path=from_path,
+      to_path=data_dir,
+      remove_finished=True
+  )
+
+  dataset_dir = pathlib.Path(os.path.join(data_dir, 'hymenoptera_data'))
+  print(f'Data is downloaded to: {dataset_dir}')
+
+  return dataset_dir
+
+def load_dataset(data_dir):
+  # Data augmentation and normalization for training
+  # Just normalization for validation
+  data_transforms = {
+      'train': transforms.Compose([
+          transforms.RandomResizedCrop(224),
+          transforms.RandomHorizontalFlip(),
+          transforms.ToTensor(),
+          transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+      ]),
+      'val': transforms.Compose([
+          transforms.Resize(256),
+          transforms.CenterCrop(224),
+          transforms.ToTensor(),
+          transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+      ]),
+  }
+
+  image_datasets = {
+      x: datasets.ImageFolder(
+          os.path.join(data_dir, x),
+          data_transforms[x]
+      )
+      for x in ['train', 'val']
+  }
+
+  class_names = image_datasets['train'].classes
+  print(f'Class names: {class_names}')
+  print(f'Number of classes: {len(class_names)}')
+
+  return image_datasets, class_names
+
+def train(model, criterion, optimizer, scheduler, dataset_sizes, dataloaders, device, epochs, writer):
+  since = time.time()
+
+  best_model_wts = copy.deepcopy(model.state_dict())
+  best_acc = 0.0
+
+  for epoch in range(epochs):
+    print('Epoch {}/{}'.format(epoch, epochs - 1))
+    print('-' * 10)
+
+    # Each epoch has a training and validation phase
+    for phase in ['train', 'val']:
+      if phase == 'train':
+        model.train()  # Set model to training mode
+      else:
+        model.eval()   # Set model to evaluate mode
+
+      running_loss = 0.0
+      running_corrects = 0
+
+      # Iterate over data.
+      for inputs, labels in dataloaders[phase]:
+        inputs = inputs.to(device)
+        labels = labels.to(device)
+
+        # zero the parameter gradients
+        optimizer.zero_grad()
+
+        # forward
+        # track history if only in train
+        with torch.set_grad_enabled(phase == 'train'):
+          outputs = model(inputs)
+          _, preds = torch.max(outputs, 1)
+          loss = criterion(outputs, labels)
+
+          # backward + optimize only if in training phase
+          if phase == 'train':
+            loss.backward()
+            optimizer.step()
+
+        # statistics
+        running_loss += loss.item() * inputs.size(0)
+        running_corrects += torch.sum(preds == labels.data)
+      if phase == 'train':
+        scheduler.step()
+
+      epoch_loss = running_loss / dataset_sizes[phase]
+      epoch_acc = running_corrects.double() / dataset_sizes[phase]
+
+      if phase == 'train':
+        writer.add_scalar('Loss/train', epoch_loss, epoch)
+        writer.add_scalar('Accuracy/train', epoch_acc, epoch)
+
+      if phase == 'val':
+        writer.add_scalar('Loss/test', epoch_loss, epoch)
+        writer.add_scalar('Accuracy/test', epoch_acc, epoch)
+
+      print('{} Loss: {:.4f} Acc: {:.4f}'.format(
+          phase, epoch_loss, epoch_acc))
+
+      # deep copy the model
+      if phase == 'val' and epoch_acc > best_acc:
+        best_acc = epoch_acc
+        best_model_wts = copy.deepcopy(model.state_dict())
+
+  time_elapsed = time.time() - since
+  print('Training complete in {:.0f}m {:.0f}s'.format(
+      time_elapsed // 60, time_elapsed % 60))
+  print('Best val Acc: {:4f}'.format(best_acc))
+
+  # load best model weights
+  model.load_state_dict(best_model_wts)
+  return model
+
+def load_model(class_names, device, pretrained=True):
+
+  model_ft = models.resnet18(pretrained=pretrained)
+  num_ftrs = model_ft.fc.in_features
+  model_ft.fc = torch.nn.Linear(num_ftrs, len(class_names))
+  model_ft = model_ft.to(device)
+
+  return model_ft
+
+def main():
+
+  args = parse_args()
+
+  local_data_dir = './tmp/data'
+
+  local_model_dir = './tmp/model'
+  local_tensorboard_log_dir = './tmp/logs'
+
+  #TODO: update when gcsfuse ready
+  gcsfuse_ready = False
+
+  model_dir = (gcsfuse_ready and args.model_dir) or local_model_dir
+  tensorboard_log_dir = (gcsfuse_ready and
+                         args.tensorboard_log_dir) or local_tensorboard_log_dir
+
+  writer = SummaryWriter(tensorboard_log_dir)
+
+  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+  print(f'Device: {device}')
+
+  data_dir = download_data(local_data_dir)
+  image_datasets, class_names = load_dataset(data_dir)
+
+  dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
+  print(f'Dataset sizes: {dataset_sizes}')
+
+  if model_dir == local_model_dir:
+    utils.makedirs(model_dir)
+
+  dataloaders = {
+      x: torch.utils.data.DataLoader(
+          image_datasets[x],
+          batch_size=args.batch_size,
+          shuffle=True,
+          num_workers=args.num_workers
+      )
+      for x in ['train', 'val']
+  }
+
+  model_ft = load_model(class_names, device)
+
+  criterion = torch.nn.CrossEntropyLoss()
+
+  # Observe that all parameters are being optimized
+  optimizer_ft = optim.SGD(
+      model_ft.parameters(), lr=args.learning_rate, momentum=args.momentum)
+
+  # Decay LR by a factor of 0.1 every 7 epochs
+  exp_lr_scheduler = optim.lr_scheduler.StepLR(
+      optimizer_ft, step_size=7, gamma=0.1)
+
+  model = train(
+      model=model_ft,
+      criterion=criterion,
+      optimizer=optimizer_ft,
+      scheduler=exp_lr_scheduler,
+      dataset_sizes=dataset_sizes,
+      dataloaders=dataloaders,
+      device=device,
+      epochs=args.epochs,
+      writer=writer,
+  )
+  model_name = 'antandbee.pth'
+  model_path = os.path.join(model_dir, f'{model_name}')
+
+  torch.save(model.state_dict(), model_path)
+  print(f'Model is saved to {model_dir}')
+  utils.gcs_upload(
+      dir=model_dir,
+      local_dir=local_model_dir,
+      gcs_dir=args.model_dir,
+      gcsfuse_ready=gcsfuse_ready,
+      local_mode=args.local_mode
+  )
+
+  print(f'Tensorboard logs are saved to: {tensorboard_log_dir}')
+  utils.gcs_upload(
+      dir=tensorboard_log_dir,
+      local_dir=local_tensorboard_log_dir,
+      gcs_dir=args.tensorboard_log_dir,
+      gcsfuse_ready=gcsfuse_ready,
+      local_mode=args.local_mode
+  )
+
+  writer.close()
+
+  return
+
+if __name__ == '__main__':
+  main()
@@ -0,0 +1,31 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the \"License\");
+# you may not use this file except in compliance with the License.\n",
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License 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.
+
+import os
+import shutil
+import subprocess
+
+def makedirs(model_dir):
+  if os.path.exists(model_dir) and os.path.isdir(model_dir):
+    shutil.rmtree(model_dir)
+  os.makedirs(model_dir)
+  return
+
+def gcs_upload(dir, local_dir, gcs_dir, gcsfuse_ready, local_mode):
+  if not local_mode and dir == local_dir and not gcsfuse_ready:
+    subprocess.run(['gsutil', 'cp', '-r',
+                    local_dir,
+                    os.path.dirname(gcs_dir)])
+    print(f'{local_dir} is uploaded to {gcs_dir}')
+  return
@@ -0,0 +1,720 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# Copyright 2021 Google LLC\n",
+    "#\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#\n",
+    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License."
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# Vertex Prediction with Custom TorchServe Container"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "<table align=\"left\">\n",
+    "  <td>\n",
+    "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/community-content/pytorch_image_classification_single_gpu_with_vertex_sdk_and_torchserve/vertex_prediction_with_custom_torchserve_container.ipynb\">\n",
+    "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+    "      View on GitHub\n",
+    "    </a>\n",
+    "  </td>\n",
+    "</table>"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Setup"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "PROJECT_ID = \"YOUR PROJECT ID\"\n",
+    "BUCKET_NAME = \"gs://YOUR BUCKET NAME\"\n",
+    "REGION = \"YOUR REGION\"\n",
+    "SERVICE_ACCOUNT = \"YOUR SERVICE ACCOUNT\""
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "content_name = \"pt-img-cls-gpu-cust-cont-torchserve\""
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Training Artifact"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "gcs_output_uri_prefix = f\"{BUCKET_NAME}/{content_name}\""
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! gsutil ls $gcs_output_uri_prefix"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Vertex Prediction using Custom TorchServe Container"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Test Sample Image"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! curl -O https://raw.githubusercontent.com/alvarobartt/pytorch-model-serving/master/images/sample.jpg"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! ls sample.jpg"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%run convert_b64.py"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! ls sample_b64.json"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Model Archive for TorchServe"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! gsutil cp -r $gcs_output_uri_prefix/model ./model_server/"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! ls ./model_server/model/"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! cd model_server && torch-model-archiver \\\n",
+    "     --model-name antandbee \\\n",
+    "     --version 1.0 \\\n",
+    "     --serialized-file ./model/antandbee.pth \\\n",
+    "     --model-file ./model.py \\\n",
+    "     --handler ./handler.py \\\n",
+    "     --extra-files ./index_to_name.json \\\n",
+    "     -f"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! ls model_server/antandbee.mar"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Option: TorchServe Local Run\n",
+    "\n",
+    "```\n",
+    "cd model_server\n",
+    "torchserve --model-store ./ \\\n",
+    "  --ts-config ./config.properties \\\n",
+    "  --models antandbee=antandbee.mar\n",
+    "\n",
+    "curl http://localhost:8080/ping\n",
+    "\n",
+    "curl http://127.0.0.1:8081/models/antandbee\n",
+    "\n",
+    "curl -X POST \\\n",
+    "  -H \"Content-Type: application/json; charset=utf-8\" \\\n",
+    "  -d @sample_b64.json \\\n",
+    "  http://localhost:8080/predictions/antandbee\n",
+    "\n",
+    "torchserve --stop\n",
+    "\n",
+    "\n",
+    "! rm model_server/antandbee.mar\n",
+    "! rm -rf model_server/logs\n",
+    "\n",
+    "```"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Custom TorchServe Container"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "hostname = \"gcr.io\"\n",
+    "tag = \"latest\"\n",
+    "\n",
+    "model_name = \"antandbee\"\n",
+    "image_name_serve = content_name + \"-\" + model_name\n",
+    "custom_container_image_uri_serve=f\"{hostname}/{PROJECT_ID}/{image_name_serve}:{tag}\""
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! cd model_server && docker build -t $custom_container_image_uri_serve -f Dockerfile ."
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! rm -rf ./model_server/model/"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! docker run \\\n",
+    "    --rm -it \\\n",
+    "    -d \\\n",
+    "    --name ts_antandbee \\\n",
+    "    -p 8080:8080 \\\n",
+    "    -p 8081:8081 \\\n",
+    "    $custom_container_image_uri_serve"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! curl http://localhost:8080/ping"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! curl http://127.0.0.1:8081/models/antandbee"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! curl -X POST \\\n",
+    "  -H \"Content-Type: application/json; charset=utf-8\" \\\n",
+    "  -d @sample_b64.json \\\n",
+    "  localhost:8080/predictions/antandbee"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! docker stop ts_antandbee"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! docker push $custom_container_image_uri_serve"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! gcloud container images list --repository $hostname/$PROJECT_ID"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Initialize Vertex SDK"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! pip install -r requirements.txt"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from google.cloud import aiplatform\n",
+    "\n",
+    "aiplatform.init(\n",
+    "    project=PROJECT_ID,\n",
+    "    staging_bucket=BUCKET_NAME,\n",
+    "    location=REGION,\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Create a Vertex Model with Custom TorchServe Container"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "model_display_name = image_name_serve"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "model = aiplatform.Model.upload(\n",
+    "    display_name=model_display_name,\n",
+    "    serving_container_image_uri=custom_container_image_uri_serve,\n",
+    "    serving_container_ports=[8080],\n",
+    "    serving_container_predict_route=f\"/predictions/{model_name}\",\n",
+    "    serving_container_health_route=\"/ping\",\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Create a Vertex Endpoint for Online Prediction"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "endpoint = model.deploy(\n",
+    "    machine_type=\"n1-standard-4\",\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "endpoint.resource_name"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "import base64\n",
+    "\n",
+    "def convert_b64(input_file_name):\n",
+    "  \"\"\"Open image and convert it to Base64\"\"\"\n",
+    "  with open(input_file_name, 'rb') as input_file:\n",
+    "    jpeg_bytes = base64.b64encode(input_file.read()).decode('utf-8')\n",
+    "  return jpeg_bytes"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "image_file_name = \"./sample.jpg\"\n",
+    "instance = {\"data\": {\"b64\": convert_b64(image_file_name)}}\n",
+    "prediction = endpoint.predict(instances=[instance])\n",
+    "print(prediction)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Clean Up"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! gsutil rm -rf $gcs_output_uri_prefix"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! rm sample.jpg\n",
+    "! rm sample_b64.json"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! rm model_server/antandbee.mar\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
@@ -0,0 +1,507 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# Copyright 2021 Google LLC\n",
+    "#\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#\n",
+    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License."
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# PyTorch Image Classification Single GPU using Vertex Training with Custom Container"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "<table align=\"left\">\n",
+    "  <td>\n",
+    "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/community-content/pytorch_image_classification_single_gpu_with_vertex_sdk_and_torchserve/vertex_training_with_custom_container.ipynb\">\n",
+    "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
+    "      View on GitHub\n",
+    "    </a>\n",
+    "  </td>\n",
+    "</table>"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Setup"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "PROJECT_ID = \"YOUR PROJECT ID\"\n",
+    "BUCKET_NAME = \"gs://YOUR BUCKET NAME\"\n",
+    "REGION = \"YOUR REGION\"\n",
+    "SERVICE_ACCOUNT = \"YOUR SERVICE ACCOUNT\""
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "content_name = \"pt-img-cls-gpu-cust-cont-torchserve\""
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Local Training"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! ls trainer"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! cat trainer/requirements.txt"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! pip install -r trainer/requirements.txt"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! cat trainer/task.py"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%run trainer/task.py --epochs 5 --local-mode"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! ls ./tmp"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! rm -rf ./tmp"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Vertex Training using Vertex SDK and Custom Container"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Build Custom Container"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "hostname = \"gcr.io\"\n",
+    "image_name_train = content_name\n",
+    "tag = \"latest\"\n",
+    "\n",
+    "custom_container_image_uri_train=f\"{hostname}/{PROJECT_ID}/{image_name_train}:{tag}\""
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! cd trainer && docker build -t $custom_container_image_uri_train -f Dockerfile ."
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! docker run --rm $custom_container_image_uri_train --epochs 5 --local-mode"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! docker push $custom_container_image_uri_train"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! gcloud container images list --repository $hostname/$PROJECT_ID"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Initialize Vertex SDK"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! pip install -r requirements.txt"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from google.cloud import aiplatform\n",
+    "\n",
+    "aiplatform.init(\n",
+    "    project=PROJECT_ID,\n",
+    "    staging_bucket=BUCKET_NAME,\n",
+    "    location=REGION,\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Create a Vertex Tensorboard Instance"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "tensorboard = aiplatform.Tensorboard.create(\n",
+    "    display_name=content_name,\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "#### Option: Use a Previously Created Vertex Tensorboard Instance\n",
+    "\n",
+    "```\n",
+    "tensorboard_name = \"Your Tensorboard Resource Name or Tensorboard ID\"\n",
+    "tensorboard = aiplatform.Tensorboard(tensorboard_name=tensorboard_name)\n",
+    "```"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Run a Vertex SDK CustomContainerTrainingJob"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "display_name = content_name\n",
+    "gcs_output_uri_prefix = f\"{BUCKET_NAME}/{display_name}\"\n",
+    "\n",
+    "machine_type = \"n1-standard-4\"\n",
+    "accelerator_count = 1\n",
+    "accelerator_type = \"NVIDIA_TESLA_K80\"\n",
+    "\n",
+    "container_args = [\n",
+    "    '--batch-size', '256',\n",
+    "    '--epochs', '100',\n",
+    "]"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "custom_container_training_job = aiplatform.CustomContainerTrainingJob(\n",
+    "    display_name=display_name,\n",
+    "    container_uri=custom_container_image_uri_train,\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "custom_container_training_job.run(\n",
+    "    args=container_args,\n",
+    "    base_output_dir=gcs_output_uri_prefix,\n",
+    "    machine_type=machine_type,\n",
+    "    accelerator_type=accelerator_type,\n",
+    "    accelerator_count=accelerator_count,\n",
+    "    tensorboard=tensorboard.resource_name,\n",
+    "    service_account=SERVICE_ACCOUNT,\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "print(f'Custom Training Job Name: {custom_container_training_job.resource_name}')\n",
+    "print(f'GCS Output URI Prefix: {gcs_output_uri_prefix}')"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Training Artifact"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "! gsutil ls $gcs_output_uri_prefix\n"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 2
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython2",
+   "version": "2.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
				`@@ -0,0 +1 @@`
				`{"0": ["id01", "ant"], "1": ["id02", "bee"]}`