content: add tf-keras-img-cls-dist-multi-worker-cpu-cust-cont (#39)

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/multi_worker_vertex_training_on_cpu_with_custom_container.ipynb

@@ -0,0 +1,423 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# Copyright 2021 Google LLC\n",
+    "#\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#\n",
+    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License."
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# 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"
+    }
+   }
+  },
+  {
+   "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
+}

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/multi_worker_vertex_training_on_gpu_with_custom_container.ipynb

@@ -0,0 +1,427 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# Copyright 2021 Google LLC\n",
+    "#\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#\n",
+    "#     https://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License."
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "# 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"
+    }
+   }
+  },
+  {
+   "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
+}

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/requirements.txt

@@ -0,0 +1 @@
+google-cloud-aiplatform

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/trainer/cpu.Dockerfile

@@ -0,0 +1,9 @@
+FROM gcr.io/deeplearning-platform-release/tf2-cpu.2-4
+
+COPY . /trainer
+
+WORKDIR /trainer
+
+RUN pip install -r requirements.txt
+
+ENTRYPOINT ["python", "-m", "task"]

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/trainer/distribution_utils.py

@@ -0,0 +1,74 @@
+# 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 json
+
+import tensorflow as tf
+
+def get_default_strategy():
+  devices = ["device:CPU:0"]
+  return tf.distribute.MirroredStrategy(devices)
+
+def get_distribution_multiworkermirrored_strategy():
+  return tf.distribute.experimental.MultiWorkerMirroredStrategy()
+
+def get_strategy(num_worker):
+  if num_worker == 1:
+    print("distribution_strategy: tf.distribute.MirroredStrategy")
+    return get_default_strategy()
+  print("distribution_strategy: tf.distribute.MultiWorkerMirroredStrategy")
+  return get_distribution_multiworkermirrored_strategy()
+
+def setup():
+  tf_config_string = os.getenv('TF_CONFIG', None)
+  task_type, task_id = None, None
+  if tf_config_string is None:
+    num_worker = 1
+  else:
+    tf_config = json.loads(tf_config_string)
+    num_worker = len(tf_config['cluster'].get('chief', [])) + len(
+        tf_config['cluster'].get('worker', []))
+    task_type = tf_config['task'].get('type', 'chief')
+    task_id = tf_config['task'].get('index', 0)
+  return num_worker, task_type, task_id
+
+def _is_chief(task_type, task_id):
+  # Note: there are two possible `TF_CONFIG` configuration.
+  #   1) In addition to `worker` tasks, a `chief` task type is use;
+  #      in this case, this function should be modified to
+  #      `return task_type == 'chief'`.
+  #   2) Only `worker` task type is used; in this case, worker 0 is
+  #      regarded as the chief. The implementation demonstrated here
+  #      is for this case.
+  # For the purpose of this Colab section, the `task_type is None` case
+  # is added because it is effectively run with only a single worker.
+  return (task_type == 'worker' and task_id == 0) or task_type is None
+
+def _get_temp_dir(dirpath, task_id):
+  base_dirpath = 'workertemp_' + str(task_id)
+  temp_dir = os.path.join(dirpath, base_dirpath)
+  tf.io.gfile.makedirs(temp_dir)
+  return temp_dir
+
+def write_filepath(filepath, task_type, task_id):
+  dirpath = os.path.dirname(filepath)
+  base = os.path.basename(filepath)
+  if not _is_chief(task_type, task_id):
+    dirpath = _get_temp_dir(dirpath, task_id)
+  return os.path.join(dirpath, base)
+
+def clean_up(task_type, task_id, dir):
+  if not _is_chief(task_type, task_id):
+    tf.io.gfile.rmtree(os.path.dirname(dir))

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/trainer/gpu.Dockerfile

@@ -0,0 +1,9 @@
+FROM gcr.io/deeplearning-platform-release/tf2-gpu.2-4
+
+COPY . /trainer
+
+WORKDIR /trainer
+
+RUN pip install -r requirements.txt
+
+ENTRYPOINT ["python", "-m", "task"]

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/trainer/requirements.txt

@@ -0,0 +1 @@
+tensorflow==2.4.1

community-content/tf_keras_image_classification_distributed_multi_worker_with_vertex_sdk/trainer/task.py

@@ -0,0 +1,174 @@
+# 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 argparse
+import os
+
+import numpy as np
+import tensorflow as tf
+
+import distribution_utils
+
+def parse_args():
+
+  parser = argparse.ArgumentParser()
+
+  # Using environment variables for Cloud Storage directories
+  # see more details in https://cloud.google.com/vertex-ai/docs/training/code-requirements
+  parser.add_argument(
+      '--model-dir', default=os.getenv('AIP_MODEL_DIR'), type=str,
+      help='a Cloud Storage URI of a directory intended for saving model artifacts')
+  parser.add_argument(
+      '--tensorboard-log-dir', default=os.getenv('AIP_TENSORBOARD_LOG_DIR'), type=str,
+      help='a Cloud Storage URI of a directory intended for saving TensorBoard')
+  parser.add_argument(
+      '--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(
+      '--epochs', default=25, type=int, help='number of training epochs')
+  parser.add_argument(
+      '--batch-size', default=32, type=int, help='mini-batch size')
+  parser.add_argument(
+      '--img-height', default=180, type=int, help='image height')
+  parser.add_argument(
+      '--img-width', default=180, type=int, help='image width')
+  parser.add_argument(
+      '--seed', default=123, type=int, help='seed')
+  parser.add_argument(
+      '--model-version', default=1, type=int, help='model version')
+
+  parser.add_argument(
+      '--local-mode', action='store_true', help='use local mode when running on your local machine')
+
+  args = parser.parse_args()
+
+  return args
+
+def load_dataset(batch_size):
+  (x_train, y_train), _ = tf.keras.datasets.mnist.load_data()
+  # The `x` arrays are in uint8 and have values in the [0, 255] range.
+  # You need to convert them to float32 with values in the [0, 1] range.
+  x_train = x_train / np.float32(255)
+  y_train = y_train.astype(np.int64)
+  train_dataset = tf.data.Dataset.from_tensor_slices(
+      (x_train, y_train)).shuffle(60000).batch(batch_size)
+  return train_dataset
+
+def build_model():
+  model = tf.keras.Sequential([
+      tf.keras.layers.InputLayer(input_shape=(28, 28)),
+      tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
+      tf.keras.layers.Conv2D(32, 3, activation='relu'),
+      tf.keras.layers.Flatten(),
+      tf.keras.layers.Dense(128, activation='relu'),
+      tf.keras.layers.Dense(10)
+  ])
+  model.compile(
+      loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+      optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
+      metrics=['accuracy'])
+  return model
+
+def train(
+    model,
+    train_dataset,
+    epochs,
+    tensorboard_log_dir,
+    checkpoint_dir
+):
+
+  tensorboard_callback = tf.keras.callbacks.TensorBoard(
+      log_dir=tensorboard_log_dir,
+      update_freq=1
+  )
+  backup_and_restore_callback = tf.keras.callbacks.experimental.BackupAndRestore(
+      backup_dir=checkpoint_dir
+  )
+
+  history = model.fit(
+      train_dataset,
+      epochs=epochs,
+      callbacks=[
+          tensorboard_callback,
+          backup_and_restore_callback,
+      ]
+  )
+
+  print('Training accuracy: {acc}, loss: {loss}'.format(
+      acc=history.history['accuracy'][-1],
+      loss=history.history['loss'][-1])
+  )
+
+  return
+
+def main():
+
+  args = parse_args()
+
+  local_model_dir = './tmp/model'
+  local_tensorboard_log_dir = './tmp/logs'
+  local_checkpoint_dir = './tmp/checkpoints'
+
+  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
+
+  num_worker, task_type, task_id = distribution_utils.setup()
+  print(f'task_type: {task_type}, '
+        f'task_id: {task_id}, '
+        f'num_worker: {num_worker} \n'
+        )
+
+  strategy = distribution_utils.get_strategy(num_worker=num_worker)
+
+  global_batch_size = args.batch_size * num_worker
+  print(f'Global batch size: {global_batch_size}')
+
+  train_ds = load_dataset(batch_size=global_batch_size)
+
+  if num_worker > 1:
+    options = tf.data.Options()
+    options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
+    train_ds = train_ds.with_options(options)
+
+  train_ds = train_ds.prefetch(tf.data.experimental.AUTOTUNE)
+
+  with strategy.scope():
+    model = build_model()
+    latest_ckpt = tf.train.latest_checkpoint(checkpoint_dir)
+    if latest_ckpt:
+      model.load_weights(latest_ckpt)
+
+  train(
+      model=model,
+      train_dataset=train_ds,
+      epochs=args.epochs,
+      tensorboard_log_dir=tensorboard_log_dir,
+      checkpoint_dir=checkpoint_dir
+  )
+
+  model_path = os.path.join(model_dir, str(args.model_version))
+  model_path = distribution_utils.write_filepath(model_path, task_type, task_id)
+  model.save(model_path)
+  print(f'Model version {args.model_version} is saved to {model_dir}')
+
+  distribution_utils.clean_up(task_type, task_id, model_path)
+
+  print(f'Tensorboard logs are saved to: {tensorboard_log_dir}')
+
+  return
+
+if __name__ == '__main__':
+  main()