change to gcsfuse

fix: add replica_count
2026-09-26 14:42:04 +00:00 · 2021-10-25 13:28:56 -07:00 · 2021-10-25 10:23:47 -07:00
14 changed files with 3011 additions and 3600 deletions
@@ -1,391 +1,347 @@
 {
- "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"
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a6b56b1c7b76"
+      },
+      "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."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "20a5ea0081d0"
+      },
+      "source": [
+        "# PyTorch Image Classification Multi-Node Distributed Data Parallel Training on GPU using Vertex Training with Custom Container"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8752d4a255fb"
+      },
+      "source": [
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_nccl_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>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "03d216c7f7b1"
+      },
+      "source": [
+        "## Setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c5ac73516218"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"YOUR PROJECT ID\"\n",
+        "BUCKET_NAME = \"gs://YOUR BUCKET NAME\"\n",
+        "REGION = \"YOUR REGION\"\n",
+        "SERVICE_ACCOUNT = \"YOUR SERVICE ACCOUNT\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0b5ae674177e"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil ls -al $BUCKET_NAME"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "19a9b3bdd553"
+      },
+      "outputs": [],
+      "source": [
+        "content_name = \"pt-img-cls-multi-node-ddp-cust-cont\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5307fe28b633"
+      },
+      "source": [
+        "## Vertex Training using Vertex SDK and Custom Container"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "46cb58c7fbf9"
+      },
+      "source": [
+        "### Built Custom Container"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "97e66e9f9bab"
+      },
+      "outputs": [],
+      "source": [
+        "hostname = \"gcr.io\"\n",
+        "image_name = content_name\n",
+        "tag = \"latest\"\n",
+        "\n",
+        "custom_container_image_uri = f\"{hostname}/{PROJECT_ID}/{image_name}:{tag}\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ae9b29c4773f"
+      },
+      "source": [
+        "### Initialize Vertex SDK"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "dc1e84d5dec2"
+      },
+      "outputs": [],
+      "source": [
+        "! pip install -r requirements.txt"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6964be27b98e"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    staging_bucket=BUCKET_NAME,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "594a91f438f2"
+      },
+      "source": [
+        "### Create a Vertex Tensorboard Instance"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "93134273261e"
+      },
+      "outputs": [],
+      "source": [
+        "content_name = content_name + \"-gpu\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c2bd82dbcd9b"
+      },
+      "outputs": [],
+      "source": [
+        "tensorboard = aiplatform.Tensorboard.create(\n",
+        "    display_name=content_name,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ebc593c6472e"
+      },
+      "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",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "0769e8e34c2f"
+      },
+      "source": [
+        "### Run a Vertex SDK CustomContainerTrainingJob"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "023f33ece826"
+      },
+      "outputs": [],
+      "source": [
+        "display_name = content_name\n",
+        "gcs_output_uri_prefix = f\"{BUCKET_NAME}/{display_name}\"\n",
+        "\n",
+        "replica_count = 1\n",
+        "machine_type = \"n1-standard-4\"\n",
+        "accelerator_count = 4\n",
+        "accelerator_type = \"NVIDIA_TESLA_K80\"\n",
+        "\n",
+        "args = [\n",
+        "    \"--backend\",\n",
+        "    \"nccl\",\n",
+        "    \"--batch-size\",\n",
+        "    \"128\",\n",
+        "    \"--epochs\",\n",
+        "    \"25\",\n",
+        "]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d4b599e726ef"
+      },
+      "outputs": [],
+      "source": [
+        "custom_container_training_job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=display_name,\n",
+        "    container_uri=custom_container_image_uri,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "81321e3bdf7f"
+      },
+      "outputs": [],
+      "source": [
+        "custom_container_training_job.run(\n",
+        "    args=args,\n",
+        "    base_output_dir=gcs_output_uri_prefix,\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    accelerator_count=accelerator_count,\n",
+        "    accelerator_type=accelerator_type,\n",
+        "    tensorboard=tensorboard.resource_name,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "5100712c2c4c"
+      },
+      "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}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f9b77676e5a6"
+      },
+      "source": [
+        "### Training Output Artifact"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0e171ce95ace"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil ls $gcs_output_uri_prefix"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cf1b74a12b87"
+      },
+      "source": [
+        "## Clean Up Artifact"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a0b15089c341"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil rm -rf $gcs_output_uri_prefix"
+      ]
    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "# PyTorch Image Classification Multi-Node Distributed Data Parallel Training on 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_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_nccl_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"
+  ],
+  "metadata": {
+    "colab": {
+      "name": "multi_node_ddp_nccl_vertex_training_with_custom_container.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
    }
-   }
  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! gsutil ls -al $BUCKET_NAME"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "content_name = \"pt-img-cls-multi-node-ddp-cust-cont\""
-   ],
-   "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": [
-    "### Built Custom Container"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "hostname = \"gcr.io\"\n",
-    "image_name = content_name\n",
-    "tag = \"latest\"\n",
-    "\n",
-    "custom_container_image_uri=f\"{hostname}/{PROJECT_ID}/{image_name}:{tag}\""
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### Initialize Vertex SDK"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! pip install -r requirements.txt"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "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"
-    }
-   },
-   "execution_count": null,
-   "outputs": []
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### Create a Vertex Tensorboard Instance"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "content_name = content_name + \"-gpu\"\n",
-    "\n",
-    "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",
-    "replica_count = 4\n",
-    "machine_type = \"n1-standard-4\"\n",
-    "accelerator_count = 1\n",
-    "accelerator_type = \"NVIDIA_TESLA_K80\"\n",
-    "\n",
-    "args = [\n",
-    "    '--backend', 'nccl',\n",
-    "    '--batch-size', '128',\n",
-    "    '--epochs', '25',\n",
-    "]"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "source": [
-    "custom_container_training_job = aiplatform.CustomContainerTrainingJob(\n",
-    "    display_name=display_name,\n",
-    "    container_uri=custom_container_image_uri,\n",
-    ")"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
-   "execution_count": null,
-   "outputs": []
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "custom_container_training_job.run(\n",
-    "    args=args,\n",
-    "    base_output_dir=gcs_output_uri_prefix,\n",
-    "    machine_type=machine_type,\n",
-    "    accelerator_count=accelerator_count,\n",
-    "    accelerator_type=accelerator_type,\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 Output Artifact"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! gsutil ls $gcs_output_uri_prefix"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## Clean Up Artifact"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! gsutil rm -rf $gcs_output_uri_prefix"
-   ],
-   "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
-}
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -19,7 +19,7 @@ Adapted from: https://github.com/narumiruna/pytorch-distributed-example

 import argparse
 import os
-import tempfile
+import shutil

 import torch
 from torch import distributed
@@ -29,8 +29,6 @@ from torch.utils.tensorboard import SummaryWriter

 from torchvision import datasets, transforms

-import utils
-
 def parse_args():

  parser = argparse.ArgumentParser()
@@ -43,6 +41,9 @@ def parse_args():
  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(
+      '--checkpoint-dir', default=os.getenv('AIP_CHECKPOINT_DIR'), type=str,
+      help='a Cloud Storage URI of a directory intended for saving checkpoints')

  parser.add_argument(
      '--backend', type=str, default='gloo',
@@ -68,7 +69,6 @@ def parse_args():
      '-lr', '--learning-rate', type=float, default=1e-3)
  parser.add_argument(
      '--batch-size', type=int, default=128)
-
  parser.add_argument(
      '--local-mode', action='store_true', help='use local mode when running on your local machine')

@@ -76,6 +76,12 @@ def parse_args():

  return args

+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 distributed_is_initialized():
  if distributed.is_available():
    if distributed.is_initialized():
@@ -176,7 +182,7 @@ class Trainer(object):

    for epoch in range(1, epochs + 1):

-      print("Epoch: {}, Training ...".format(epoch))
+      print('Epoch: {}, Training ...'.format(epoch))
      train_loss, train_acc = self.train()

      if is_chief:
@@ -242,16 +248,30 @@ def main():
  local_data_dir = './tmp/data'
  local_model_dir = './tmp/model'
  local_tensorboard_log_dir = './tmp/logs'
+  local_checkpoint_dir = './tmp/checkpoints'

-  #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
+  model_dir = args.model_dir or local_model_dir
+  tensorboard_log_dir = args.tensorboard_log_dir or local_tensorboard_log_dir
+  checkpoint_dir = args.checkpoint_dir or local_checkpoint_dir
+
+  gs_prefix = 'gs://'
+  gcsfuse_prefix = '/gcs/'
+  if model_dir and model_dir.startswith(gs_prefix):
+    model_dir = model_dir.replace(gs_prefix, gcsfuse_prefix)
+  if tensorboard_log_dir and tensorboard_log_dir.startswith(gs_prefix):
+    tensorboard_log_dir = tensorboard_log_dir.replace(gs_prefix, gcsfuse_prefix)
+  if checkpoint_dir and checkpoint_dir.startswith(gs_prefix):
+    checkpoint_dir = checkpoint_dir.replace(gs_prefix, gcsfuse_prefix)

  writer = SummaryWriter(tensorboard_log_dir)

  is_chief = args.rank == 0
+  if is_chief:
+    makedirs(checkpoint_dir)
+    print(f'Checkpoints will be saved to {checkpoint_dir}')
+
+  checkpoint_path = os.path.join(checkpoint_dir, 'checkpoint.pt')
+  print(f'checkpoint_path is {checkpoint_path}')

  if args.world_size > 1:
    print('Initializing distributed backend with {} nodes'.format(args.world_size))
@@ -261,36 +281,38 @@ def main():
          world_size=args.world_size,
          rank=args.rank,
      )
-    print(f"[{os.getpid()}]: "
-          f"world_size = {distributed.get_world_size()}, "
-          f"rank = {distributed.get_rank()}, "
-          f"backend={distributed.get_backend()} \n", end='')
+    print(f'[{os.getpid()}]: '
+          f'world_size = {distributed.get_world_size()}, '
+          f'rank = {distributed.get_rank()}, '
+          f'backend={distributed.get_backend()} \n', end='')

  if torch.cuda.is_available() and not args.no_cuda:
-    device = torch.device("cuda:{}".format(args.rank))
+    device = torch.device('cuda:{}'.format(args.rank))
  else:
-    device = torch.device("cpu")
+    device = torch.device('cpu')

  model = Net(device=device)
  if distributed_is_initialized():
+    model.to(device)
    model = DistributedDataParallel(model)

-  checkpoint_path = tempfile.gettempdir() + "/model.checkpoint"
  if is_chief:
    # All processes should see same parameters as they all start from same
    # random parameters and gradients are synchronized in backward passes.
    # Therefore, saving it in one process is sufficient.
    torch.save(model.state_dict(), checkpoint_path)
+    print(f'Initial chief checkpoint is saved to {checkpoint_path}')

  # Use a barrier() to make sure that process 1 loads the model after process
  # 0 saves it.
  if distributed_is_initialized():
    distributed.barrier()
    # configure map_location properly
-    map_location = {'cuda:%d' % 0: 'cuda:%d' % args.rank}
-    model.load_state_dict(torch.load(checkpoint_path, map_location=map_location))
+    model.load_state_dict(torch.load(checkpoint_path, map_location=device))
+    print(f'Initial chief checkpoint is saved to {checkpoint_path} with map_location {device}')
  else:
    model.load_state_dict(torch.load(checkpoint_path))
+    print(f'Initial chief checkpoint is loaded from {checkpoint_path}')

  optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)

@@ -310,29 +332,12 @@ def main():
  )
  trainer.fit(args.epochs, is_chief, writer)

-  if is_chief and model_dir == local_model_dir:
-    utils.makedirs(model_dir)
+  if model_dir == local_model_dir:
+    makedirs(model_dir)
    trainer.save(model_dir)
    print(f'Model is saved to {model_dir}')

-  if is_chief and not args.local_mode:
-    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}')
-  if is_chief and not args.local_mode:
-    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()

@@ -1,31 +0,0 @@
-# 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
@@ -16,6 +16,7 @@ import argparse
 import copy
 import os
 import pathlib
+import shutil
 import time

 import torch
@@ -24,8 +25,6 @@ from torch.utils.tensorboard import SummaryWriter

 from torchvision import datasets, models, transforms

-import utils
-
 def parse_args():

  parser = argparse.ArgumentParser()
@@ -62,6 +61,12 @@ def parse_args():

  return args

+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 download_data(data_dir):

  dataset_url = 'https://download.pytorch.org/tutorial/hymenoptera_data.zip'
@@ -207,13 +212,17 @@ def main():
  local_model_dir = './tmp/model'
  local_tensorboard_log_dir = './tmp/logs'

-  #TODO: update when gcsfuse ready
-  gcsfuse_ready = False
+  model_dir = args.model_dir or local_model_dir
+  tensorboard_log_dir = args.tensorboard_log_dir or local_tensorboard_log_dir

-  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
+  gs_prefix = 'gs://'
+  gcsfuse_prefix = '/gcs/'
+  if model_dir and model_dir.startswith(gs_prefix):
+    model_dir = model_dir.replace(gs_prefix, gcsfuse_prefix)
+  if tensorboard_log_dir and tensorboard_log_dir.startswith(gs_prefix):
+    tensorboard_log_dir = tensorboard_log_dir.replace(gs_prefix, gcsfuse_prefix)

+  makedirs(model_dir)
  writer = SummaryWriter(tensorboard_log_dir)

  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
@@ -225,9 +234,6 @@ def main():
  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],
@@ -266,22 +272,8 @@ def main():

  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()

@@ -1,31 +0,0 @@
-# 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
@@ -1,507 +1,431 @@
 {
- "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"
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a6b56b1c7b76"
+      },
+      "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."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d229c51d7335"
+      },
+      "source": [
+        "# PyTorch Image Classification Single GPU using Vertex Training with Custom Container"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "08b6ccb2de0f"
+      },
+      "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>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "03d216c7f7b1"
+      },
+      "source": [
+        "## Setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c5ac73516218"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"YOUR PROJECT ID\"\n",
+        "BUCKET_NAME = \"gs://YOUR BUCKET NAME\"\n",
+        "REGION = \"YOUR REGION\"\n",
+        "SERVICE_ACCOUNT = \"YOUR SERVICE ACCOUNT\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "410ef4707617"
+      },
+      "outputs": [],
+      "source": [
+        "content_name = \"pt-img-cls-gpu-cust-cont-torchserve\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ed102448ae75"
+      },
+      "source": [
+        "## Local Training"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "47dfe4e1dfe6"
+      },
+      "outputs": [],
+      "source": [
+        "! ls trainer"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "be6fcb587414"
+      },
+      "outputs": [],
+      "source": [
+        "! cat trainer/requirements.txt"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "060178a8ad71"
+      },
+      "outputs": [],
+      "source": [
+        "! pip install -r trainer/requirements.txt"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4421e56c4ee3"
+      },
+      "outputs": [],
+      "source": [
+        "! cat trainer/task.py"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4bb2f6fc75cc"
+      },
+      "outputs": [],
+      "source": [
+        "%run trainer/task.py --epochs 5 --local-mode"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cb748af8f2b8"
+      },
+      "outputs": [],
+      "source": [
+        "! ls ./tmp"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "2ca1626c5dbd"
+      },
+      "outputs": [],
+      "source": [
+        "! rm -rf ./tmp"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "4eb55916219c"
+      },
+      "source": [
+        "## Vertex Training using Vertex SDK and Custom Container"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e61fc5c1b9e0"
+      },
+      "source": [
+        "### Build Custom Container"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c70a580cc6d4"
+      },
+      "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}\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "e58c75158872"
+      },
+      "outputs": [],
+      "source": [
+        "! cd trainer && docker build -t $custom_container_image_uri_train -f Dockerfile ."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bde55966086a"
+      },
+      "outputs": [],
+      "source": [
+        "! docker run --rm $custom_container_image_uri_train --epochs 5 --local-mode"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f4a3cdb78045"
+      },
+      "outputs": [],
+      "source": [
+        "! docker push $custom_container_image_uri_train"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "526115eabf35"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud container images list --repository $hostname/$PROJECT_ID"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f026b57d265e"
+      },
+      "source": [
+        "### Initialize Vertex SDK"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "19c02e9e9b8d"
+      },
+      "outputs": [],
+      "source": [
+        "! pip install -r requirements.txt"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a0deece1086e"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    staging_bucket=BUCKET_NAME,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1ec68e279ed1"
+      },
+      "source": [
+        "### Create a Vertex Tensorboard Instance"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d945f0e32b02"
+      },
+      "outputs": [],
+      "source": [
+        "tensorboard = aiplatform.Tensorboard.create(\n",
+        "    display_name=content_name,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cb13cc82acc5"
+      },
+      "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",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d04453c214d2"
+      },
+      "source": [
+        "### Run a Vertex SDK CustomContainerTrainingJob"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4b5b5dea03c8"
+      },
+      "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\",\n",
+        "    \"256\",\n",
+        "    \"--epochs\",\n",
+        "    \"100\",\n",
+        "]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b2073fb5824e"
+      },
+      "outputs": [],
+      "source": [
+        "custom_container_training_job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=display_name,\n",
+        "    container_uri=custom_container_image_uri_train,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d442a02878d1"
+      },
+      "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",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "29b14f2289f3"
+      },
+      "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}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "3ab127e5f2b2"
+      },
+      "source": [
+        "### Training Artifact"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8e0c74a00dee"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil ls $gcs_output_uri_prefix"
+      ]
    }
-   }
-  },
-  {
-   "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"
+  ],
+  "metadata": {
+    "colab": {
+      "name": "vertex_training_with_custom_container.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
    }
-   }
  },
-  {
-   "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
-}
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -1,423 +1,363 @@
 {
- "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"
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a6b56b1c7b76"
+      },
+      "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."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "ac14831de120"
+      },
+      "source": [
+        "# TF-Keras Image Classification Distributed Multi-Worker Training on CPU using Vertex Training with Custom Container"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "1f2615ea7930"
+      },
+      "source": [
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/multi_worker_vertex_training_on_cpu_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>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "03d216c7f7b1"
+      },
+      "source": [
+        "## Setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c5ac73516218"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"YOUR PROJECT ID\"\n",
+        "BUCKET_NAME = \"gs://YOUR BUCKET NAME\"\n",
+        "REGION = \"YOUR REGION\"\n",
+        "SERVICE_ACCOUNT = \"YOUR SERVICE ACCOUNT\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b9b9229d7bc4"
+      },
+      "outputs": [],
+      "source": [
+        "content_name = \"tf-keras-img-cls-dist-multi-worker-cpu-cust-cont\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "070e602ded80"
+      },
+      "source": [
+        "## Vertex Training using Vertex SDK and Custom Container"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99b9c8546a46"
+      },
+      "source": [
+        "### Build Custom Container"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f5196e294db6"
+      },
+      "outputs": [],
+      "source": [
+        "hostname = \"gcr.io\"\n",
+        "image_name = content_name\n",
+        "tag = \"latest\"\n",
+        "\n",
+        "custom_container_image_uri = f\"{hostname}/{PROJECT_ID}/{image_name}:{tag}\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "723478af31e6"
+      },
+      "outputs": [],
+      "source": [
+        "! cd trainer && docker build -t $custom_container_image_uri -f cpu.Dockerfile ."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "66ac893a871b"
+      },
+      "outputs": [],
+      "source": [
+        "! docker run --rm $custom_container_image_uri --epochs 2 --local-mode"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d131ee61bfca"
+      },
+      "outputs": [],
+      "source": [
+        "! docker push $custom_container_image_uri"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bc5682081ee8"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud container images list --repository $hostname/$PROJECT_ID"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8b9c18d51ddf"
+      },
+      "source": [
+        "### Initialize Vertex SDK"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cfee2f165779"
+      },
+      "outputs": [],
+      "source": [
+        "! pip install -r requirements.txt"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b542a33f1782"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    staging_bucket=BUCKET_NAME,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a0f25261a58c"
+      },
+      "source": [
+        "### Create a Vertex Tensorboard Instance"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "abacc13094c9"
+      },
+      "outputs": [],
+      "source": [
+        "tensorboard = aiplatform.Tensorboard.create(\n",
+        "    display_name=content_name,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "77ca3177b71e"
+      },
+      "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",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f1733451a790"
+      },
+      "source": [
+        "### Run a Vertex SDK CustomContainerTrainingJob"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0d5e38fade0a"
+      },
+      "outputs": [],
+      "source": [
+        "display_name = content_name\n",
+        "gcs_output_uri_prefix = f\"{BUCKET_NAME}/{display_name}\"\n",
+        "\n",
+        "replica_count = 4\n",
+        "machine_type = \"n1-standard-4\"\n",
+        "\n",
+        "container_args = [\n",
+        "    \"--epochs\",\n",
+        "    \"50\",\n",
+        "    \"--batch-size\",\n",
+        "    \"32\",\n",
+        "]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "23ceccc746f5"
+      },
+      "outputs": [],
+      "source": [
+        "custom_container_training_job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=display_name,\n",
+        "    container_uri=custom_container_image_uri,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "469599f8b676"
+      },
+      "outputs": [],
+      "source": [
+        "custom_container_training_job.run(\n",
+        "    args=container_args,\n",
+        "    base_output_dir=gcs_output_uri_prefix,\n",
+        "    replica_count=replica_count,\n",
+        "    machine_type=machine_type,\n",
+        "    tensorboard=tensorboard.resource_name,\n",
+        "    service_account=SERVICE_ACCOUNT,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "e35aa8a39a11"
+      },
+      "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}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "40118a4bb0de"
+      },
+      "source": [
+        "### Training Output Artifact"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6d414dfc4585"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil ls $gcs_output_uri_prefix"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "77cd165390f9"
+      },
+      "source": [
+        "## Clean Up Artifact"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9b73b8cb646d"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil rm -rf $gcs_output_uri_prefix"
+      ]
    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "# TF-Keras Image Classification Distributed Multi-Worker Training on CPU 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/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/multi_worker_vertex_training_on_cpu_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"
+  ],
+  "metadata": {
+    "colab": {
+      "name": "multi_worker_vertex_training_on_cpu_with_custom_container.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
    }
-   }
  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "content_name = \"tf-keras-img-cls-dist-multi-worker-cpu-cust-cont\""
-   ],
-   "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 = content_name\n",
-    "tag = \"latest\"\n",
-    "\n",
-    "custom_container_image_uri=f\"{hostname}/{PROJECT_ID}/{image_name}:{tag}\""
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! cd trainer && docker build -t $custom_container_image_uri -f cpu.Dockerfile ."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! docker run --rm $custom_container_image_uri --epochs 2 --local-mode"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! docker push $custom_container_image_uri"
-   ],
-   "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,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "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",
-    "replica_count = 4\n",
-    "machine_type = \"n1-standard-4\"\n",
-    "\n",
-    "container_args = [\n",
-    "    '--epochs', '50',\n",
-    "    '--batch-size', '32',\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,\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",
-    "    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 Output Artifact"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! gsutil ls $gcs_output_uri_prefix"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## Clean Up Artifact"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! gsutil rm -rf $gcs_output_uri_prefix"
-   ],
-   "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
-}
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -1,427 +1,367 @@
 {
- "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"
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a6b56b1c7b76"
+      },
+      "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."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "e8fc31cbf9f6"
+      },
+      "source": [
+        "# TF-Keras Image Classification Distributed Multi-Worker Training on GPU using Vertex Training with Custom Container"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5825c076e0f8"
+      },
+      "source": [
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/multi_worker_vertex_training_on_gpu_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>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "03d216c7f7b1"
+      },
+      "source": [
+        "## Setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c5ac73516218"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"YOUR PROJECT ID\"\n",
+        "BUCKET_NAME = \"gs://YOUR BUCKET NAME\"\n",
+        "REGION = \"YOUR REGION\"\n",
+        "SERVICE_ACCOUNT = \"YOUR SERVICE ACCOUNT\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a5f0c3700b8e"
+      },
+      "outputs": [],
+      "source": [
+        "content_name = \"tf-keras-img-cls-dist-multi-worker-gpu-cust-cont\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "070e602ded80"
+      },
+      "source": [
+        "## Vertex Training using Vertex SDK and Custom Container"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "99b9c8546a46"
+      },
+      "source": [
+        "### Build Custom Container"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "f5196e294db6"
+      },
+      "outputs": [],
+      "source": [
+        "hostname = \"gcr.io\"\n",
+        "image_name = content_name\n",
+        "tag = \"latest\"\n",
+        "\n",
+        "custom_container_image_uri = f\"{hostname}/{PROJECT_ID}/{image_name}:{tag}\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c24b9a97df06"
+      },
+      "outputs": [],
+      "source": [
+        "! cd trainer && docker build -t $custom_container_image_uri -f gpu.Dockerfile ."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "66ac893a871b"
+      },
+      "outputs": [],
+      "source": [
+        "! docker run --rm $custom_container_image_uri --epochs 2 --local-mode"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d131ee61bfca"
+      },
+      "outputs": [],
+      "source": [
+        "! docker push $custom_container_image_uri"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "bc5682081ee8"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud container images list --repository $hostname/$PROJECT_ID"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "8b9c18d51ddf"
+      },
+      "source": [
+        "### Initialize Vertex SDK"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cfee2f165779"
+      },
+      "outputs": [],
+      "source": [
+        "! pip install -r requirements.txt"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b542a33f1782"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    staging_bucket=BUCKET_NAME,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a0f25261a58c"
+      },
+      "source": [
+        "### Create a Vertex Tensorboard Instance"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "abacc13094c9"
+      },
+      "outputs": [],
+      "source": [
+        "tensorboard = aiplatform.Tensorboard.create(\n",
+        "    display_name=content_name,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "77ca3177b71e"
+      },
+      "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",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f1733451a790"
+      },
+      "source": [
+        "### Run a Vertex SDK CustomContainerTrainingJob"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3e29b50a65d3"
+      },
+      "outputs": [],
+      "source": [
+        "display_name = content_name\n",
+        "gcs_output_uri_prefix = f\"{BUCKET_NAME}/{display_name}\"\n",
+        "\n",
+        "replica_count = 4\n",
+        "machine_type = \"n1-standard-4\"\n",
+        "accelerator_count = 1\n",
+        "accelerator_type = \"NVIDIA_TESLA_K80\"\n",
+        "\n",
+        "container_args = [\n",
+        "    \"--epochs\",\n",
+        "    \"50\",\n",
+        "    \"--batch-size\",\n",
+        "    \"32\",\n",
+        "]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "23ceccc746f5"
+      },
+      "outputs": [],
+      "source": [
+        "custom_container_training_job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=display_name,\n",
+        "    container_uri=custom_container_image_uri,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "239e58fe141d"
+      },
+      "outputs": [],
+      "source": [
+        "custom_container_training_job.run(\n",
+        "    args=container_args,\n",
+        "    base_output_dir=gcs_output_uri_prefix,\n",
+        "    replica_count=replica_count,\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",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "e35aa8a39a11"
+      },
+      "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}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "40118a4bb0de"
+      },
+      "source": [
+        "### Training Output Artifact"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6d414dfc4585"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil ls $gcs_output_uri_prefix"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "77cd165390f9"
+      },
+      "source": [
+        "## Clean Up Artifact"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9b73b8cb646d"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil rm -rf $gcs_output_uri_prefix"
+      ]
    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "# TF-Keras Image Classification Distributed Multi-Worker Training on 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/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/multi_worker_vertex_training_on_gpu_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"
+  ],
+  "metadata": {
+    "colab": {
+      "name": "multi_worker_vertex_training_on_gpu_with_custom_container.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
    }
-   }
  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "content_name = \"tf-keras-img-cls-dist-multi-worker-gpu-cust-cont\""
-   ],
-   "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 = content_name\n",
-    "tag = \"latest\"\n",
-    "\n",
-    "custom_container_image_uri=f\"{hostname}/{PROJECT_ID}/{image_name}:{tag}\""
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! cd trainer && docker build -t $custom_container_image_uri -f gpu.Dockerfile ."
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! docker run --rm $custom_container_image_uri --epochs 2 --local-mode"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! docker push $custom_container_image_uri"
-   ],
-   "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,
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   }
-  },
-  {
-   "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",
-    "replica_count = 4\n",
-    "machine_type = \"n1-standard-4\"\n",
-    "accelerator_count = 1\n",
-    "accelerator_type = \"NVIDIA_TESLA_K80\"\n",
-    "\n",
-    "container_args = [\n",
-    "    '--epochs', '50',\n",
-    "    '--batch-size', '32',\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,\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 Output Artifact"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! gsutil ls $gcs_output_uri_prefix"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## Clean Up Artifact"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! gsutil rm -rf $gcs_output_uri_prefix"
-   ],
-   "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
-}
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
@@ -123,7 +123,16 @@ def main():

  model_dir = args.model_dir or local_model_dir
  tensorboard_log_dir = args.tensorboard_log_dir or local_tensorboard_log_dir
-  checkpoint_dir = local_checkpoint_dir
+  checkpoint_dir = args.checkpoint_dir or local_checkpoint_dir
+
+  gs_prefix = 'gs://'
+  gcsfuse_prefix = '/gcs/'
+  if model_dir and model_dir.startswith(gs_prefix):
+    model_dir = model_dir.replace(gs_prefix, gcsfuse_prefix)
+  if tensorboard_log_dir and tensorboard_log_dir.startswith(gs_prefix):
+    tensorboard_log_dir = tensorboard_log_dir.replace(gs_prefix, gcsfuse_prefix)
+  if checkpoint_dir and checkpoint_dir.startswith(gs_prefix):
+    checkpoint_dir = checkpoint_dir.replace(gs_prefix, gcsfuse_prefix)

  num_worker, task_type, task_id = distribution_utils.setup()
  print(f'task_type: {task_type}, '
@@ -21,7 +21,6 @@ from tensorflow.keras import layers, losses
 from tensorflow.keras.layers.experimental.preprocessing import TextVectorization

 import distribution_utils
-import utils

 VOCAB_SIZE = 10000
 MAX_SEQUENCE_LENGTH = 250
@@ -175,13 +174,18 @@ def main():
  local_checkpoint_dir = './tmp/checkpoints'
  local_tensorboard_log_dir = './tmp/logs'

-  #TODO: update when gcsfuse ready
-  gcsfuse_ready = False
-
  model_dir = args.model_dir or local_model_dir
-  checkpoint_dir = (gcsfuse_ready and
-                         args.checkpoint_dir) or local_checkpoint_dir
  tensorboard_log_dir = args.tensorboard_log_dir or local_tensorboard_log_dir
+  checkpoint_dir = args.checkpoint_dir or local_checkpoint_dir
+
+  gs_prefix = 'gs://'
+  gcsfuse_prefix = '/gcs/'
+  if model_dir and model_dir.startswith(gs_prefix):
+    model_dir = model_dir.replace(gs_prefix, gcsfuse_prefix)
+  if tensorboard_log_dir and tensorboard_log_dir.startswith(gs_prefix):
+    tensorboard_log_dir = tensorboard_log_dir.replace(gs_prefix, gcsfuse_prefix)
+  if checkpoint_dir and checkpoint_dir.startswith(gs_prefix):
+    checkpoint_dir = checkpoint_dir.replace(gs_prefix, gcsfuse_prefix)

  class_names = ['csharp', 'java', 'javascript', 'python']
  class_indices = dict(zip(class_names, range(len(class_names))))
@@ -278,13 +282,6 @@ def main():
  print(f'Tensorboard logs are saved to: {tensorboard_log_dir}')

  print(f'Checkpoints are saved to: {checkpoint_dir}')
-  utils.gcs_upload(
-      dir=checkpoint_dir,
-      local_dir=local_checkpoint_dir,
-      gcs_dir=args.checkpoint_dir,
-      gcsfuse_ready=gcsfuse_ready,
-      local_mode=args.local_mode
-  )

  return

@@ -1,31 +0,0 @@
-# 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.
-# You may obtain a copy of the License at
-#
-#     https://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
@@ -1,437 +1,380 @@
 {
- "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"
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "a6b56b1c7b76"
+      },
+      "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."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "745ff8aea147"
+      },
+      "source": [
+        "# TF-Keras Text Classification Distributed Single Worker GPUs using Vertex Training with Local Mode Container"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "a20c6440f341"
+      },
+      "source": [
+        "<table align=\"left\">\n",
+        "  <td>\n",
+        "    <a href=\"https://github.com/GoogleCloudPlatform/vertex-ai-samples/blob/master/community-content/tf_keras_text_classification_distributed_single_worker_gpus_with_gcloud_local_run_and_vertex_sdk/vertex_training_with_local_mode_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>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "03d216c7f7b1"
+      },
+      "source": [
+        "## Setup"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "c5ac73516218"
+      },
+      "outputs": [],
+      "source": [
+        "PROJECT_ID = \"YOUR PROJECT ID\"\n",
+        "BUCKET_NAME = \"gs://YOUR BUCKET NAME\"\n",
+        "REGION = \"YOUR REGION\"\n",
+        "SERVICE_ACCOUNT = \"YOUR SERVICE ACCOUNT\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "1c71d07625f5"
+      },
+      "outputs": [],
+      "source": [
+        "content_name = \"tf-keras-txt-cls-dist-single-worker-gpus-local-mode-cont\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "378392e0adb9"
+      },
+      "source": [
+        "## Local Training with Vertex Local Mode and Auto Packaging"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "05f8e4885312"
+      },
+      "outputs": [],
+      "source": [
+        "BASE_IMAGE_URI = \"us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-5:latest\"\n",
+        "SCRIPT_PATH = \"trainer/task.py\"\n",
+        "OUTPUT_IMAGE_NAME = \"gcr.io/{}/{}:latest\".format(PROJECT_ID, content_name)\n",
+        "ARGS = \"--epochs 5 --batch-size 16 --local-mode\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b45a98212ef5"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud beta ai custom-jobs local-run \\\n",
+        "  --base-image=$BASE_IMAGE_URI \\\n",
+        "  --script=$SCRIPT_PATH \\\n",
+        "  --output-image-uri=$OUTPUT_IMAGE_NAME \\\n",
+        "  -- \\\n",
+        "  $ARGS"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cbd77cb4193c"
+      },
+      "source": [
+        "## Vertex Training using Vertex SDK and Vertex Local Mode Container"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "c8b69c5e3dd4"
+      },
+      "source": [
+        "### Container Built by Vertex Local Mode"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "4869e0d1657b"
+      },
+      "outputs": [],
+      "source": [
+        "custom_container_image_uri = OUTPUT_IMAGE_NAME"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d53296526bde"
+      },
+      "outputs": [],
+      "source": [
+        "! docker push $custom_container_image_uri"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "eb193581654a"
+      },
+      "outputs": [],
+      "source": [
+        "! gcloud container images list --repository \"gcr.io\"/$PROJECT_ID"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f5b907c91ea6"
+      },
+      "source": [
+        "### Initialize Vertex SDK"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "fbfec07eaea4"
+      },
+      "outputs": [],
+      "source": [
+        "! pip install -r requirements.txt"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8fcf24c75bcd"
+      },
+      "outputs": [],
+      "source": [
+        "from google.cloud import aiplatform\n",
+        "\n",
+        "aiplatform.init(\n",
+        "    project=PROJECT_ID,\n",
+        "    staging_bucket=BUCKET_NAME,\n",
+        "    location=REGION,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "b68b79268025"
+      },
+      "source": [
+        "### Create a Vertex Tensorboard Instance"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7577c86757ab"
+      },
+      "outputs": [],
+      "source": [
+        "tensorboard = aiplatform.Tensorboard.create(\n",
+        "    display_name=content_name,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2a327cb1722d"
+      },
+      "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",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "2e1069a8f7e2"
+      },
+      "source": [
+        "### Run a Vertex SDK CustomContainerTrainingJob"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d7f0043d3e9b"
+      },
+      "outputs": [],
+      "source": [
+        "display_name = content_name\n",
+        "gcs_output_uri_prefix = f\"{BUCKET_NAME}/{display_name}\"\n",
+        "\n",
+        "machine_type = \"n1-standard-8\"\n",
+        "accelerator_count = 4\n",
+        "accelerator_type = \"NVIDIA_TESLA_P100\"\n",
+        "\n",
+        "args = [\n",
+        "    \"--epochs\",\n",
+        "    \"100\",\n",
+        "    \"--batch-size\",\n",
+        "    \"128\",\n",
+        "    \"--num-gpus\",\n",
+        "    f\"{accelerator_count}\",\n",
+        "]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "69c88c73bb5d"
+      },
+      "outputs": [],
+      "source": [
+        "custom_container_training_job = aiplatform.CustomContainerTrainingJob(\n",
+        "    display_name=display_name,\n",
+        "    container_uri=custom_container_image_uri,\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "947dd755cdc8"
+      },
+      "outputs": [],
+      "source": [
+        "custom_container_training_job.run(\n",
+        "    args=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",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "703e297f83fe"
+      },
+      "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}\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "d459452eca40"
+      },
+      "source": [
+        "### Training Output Artifact"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "d02a5890a4af"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil ls $gcs_output_uri_prefix"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "545fe0dbe7c4"
+      },
+      "source": [
+        "## Clean Up Artifact"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "b86c0fe8028e"
+      },
+      "outputs": [],
+      "source": [
+        "! gsutil rm -rf $gcs_output_uri_prefix"
+      ]
    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "# TF-Keras Text Classification Distributed Single Worker GPUs using Vertex Training with Local Mode 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/tf_keras_text_classification_distributed_single_worker_gpus_with_gcloud_local_run_and_vertex_sdk/vertex_training_with_local_mode_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"
+  ],
+  "metadata": {
+    "colab": {
+      "name": "vertex_training_with_local_mode_container.ipynb",
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
    }
-   }
  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "content_name = \"tf-keras-txt-cls-dist-single-worker-gpus-local-mode-cont\""
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## Local Training with Vertex Local Mode and Auto Packaging"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "BASE_IMAGE_URI = 'us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-5:latest'\n",
-    "SCRIPT_PATH = 'trainer/task.py'\n",
-    "OUTPUT_IMAGE_NAME = 'gcr.io/{}/{}:latest'.format(PROJECT_ID, content_name)\n",
-    "ARGS = '--epochs 5 --batch-size 16 --local-mode'"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! gcloud beta ai custom-jobs local-run \\\n",
-    "  --base-image=$BASE_IMAGE_URI \\\n",
-    "  --script=$SCRIPT_PATH \\\n",
-    "  --output-image-uri=$OUTPUT_IMAGE_NAME \\\n",
-    "  -- \\\n",
-    "  $ARGS"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## Vertex Training using Vertex SDK and Vertex Local Mode Container"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "### Container Built by Vertex Local Mode"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "custom_container_image_uri = OUTPUT_IMAGE_NAME"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! docker push $custom_container_image_uri"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! gcloud container images list --repository \"gcr.io\"/$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-8\"\n",
-    "accelerator_count = 4\n",
-    "accelerator_type = \"NVIDIA_TESLA_P100\"\n",
-    "\n",
-    "args = [\n",
-    "    '--epochs', '100',\n",
-    "    '--batch-size', '128',\n",
-    "    '--num-gpus', f'{accelerator_count}',\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,\n",
-    ")"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "custom_container_training_job.run(\n",
-    "    args=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 Output Artifact"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! gsutil ls $gcs_output_uri_prefix"
-   ],
-   "metadata": {
-    "collapsed": false,
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   }
-  },
-  {
-   "cell_type": "markdown",
-   "source": [
-    "## Clean Up Artifact"
-   ],
-   "metadata": {
-    "collapsed": false
-   }
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [
-    "! gsutil rm -rf $gcs_output_uri_prefix"
-   ],
-   "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
-}
+  "nbformat": 4,
+  "nbformat_minor": 0
+}
Author	SHA1	Message	Date
Morgan Du	b32f53f6bd	change to gcsfuse	2021-10-25 13:28:56 -07:00
Morgan Du	7808044cb8	fix: add replica_count	2021-10-25 10:23:47 -07:00