From af764b0d9866937a3a9d079ef1b40733c3ebe0d1 Mon Sep 17 00:00:00 2001 From: Morgan Du Date: Mon, 23 Aug 2021 14:44:55 -0700 Subject: [PATCH] content: add pytorch-img-cls-multi-node-ddp-cust-cont (#37) --- ...ertex_training_with_custom_container.ipynb | 554 ++++++++++++++++++ ...ertex_training_with_custom_container.ipynb | 391 ++++++++++++ .../requirements.txt | 1 + .../trainer/Dockerfile | 16 + .../trainer/__init__.py | 0 .../trainer/requirements.txt | 3 + .../trainer/task.py | 348 +++++++++++ .../trainer/utils.py | 31 + 8 files changed, 1344 insertions(+) create mode 100644 community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_gloo_vertex_training_with_custom_container.ipynb create mode 100644 community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_nccl_vertex_training_with_custom_container.ipynb create mode 100644 community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/requirements.txt create mode 100644 community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/Dockerfile create mode 100644 community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/__init__.py create mode 100644 community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/requirements.txt create mode 100644 community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/task.py create mode 100644 community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/utils.py diff --git a/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_gloo_vertex_training_with_custom_container.ipynb b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_gloo_vertex_training_with_custom_container.ipynb new file mode 100644 index 000000000..0682c5cad --- /dev/null +++ b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_gloo_vertex_training_with_custom_container.ipynb @@ -0,0 +1,554 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "outputs": [], + "source": [ + "# Copyright 2021 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ], + "metadata": { + "collapsed": false, + "pycharm": { + "name": "#%%\n" + } + } + }, + { + "cell_type": "markdown", + "source": [ + "# PyTorch Image Classification Multi-Node Distributed Data Parallel Training on CPU using Vertex Training with Custom Container" + ], + "metadata": { + "collapsed": false + } + }, + { + "cell_type": "markdown", + "source": [ + "\n", + " \n", + "
\n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + "
" + ], + "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": [ + "! 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": [ + "## 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 --no-cuda --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 = 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 Dockerfile ." + ], + "metadata": { + "collapsed": false, + "pycharm": { + "name": "#%%\n" + } + } + }, + { + "cell_type": "code", + "execution_count": null, + "outputs": [], + "source": [ + "! docker run --rm $custom_container_image_uri --epochs 5 --no-cuda --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, + "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 + \"-cpu\"\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, + "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", + "args = [\n", + " '--backend', 'gloo',\n", + " '--no-cuda',\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", + " 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}')\n" + ], + "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 +} \ No newline at end of file diff --git a/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_nccl_vertex_training_with_custom_container.ipynb b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_nccl_vertex_training_with_custom_container.ipynb new file mode 100644 index 000000000..b6adbe251 --- /dev/null +++ b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/multi_node_ddp_nccl_vertex_training_with_custom_container.ipynb @@ -0,0 +1,391 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "outputs": [], + "source": [ + "# Copyright 2021 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# https://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ], + "metadata": { + "collapsed": false, + "pycharm": { + "name": "#%%\n" + } + } + }, + { + "cell_type": "markdown", + "source": [ + "# PyTorch Image Classification Multi-Node Distributed Data Parallel Training on GPU using Vertex Training with Custom Container" + ], + "metadata": { + "collapsed": false + } + }, + { + "cell_type": "markdown", + "source": [ + "\n", + " \n", + "
\n", + " \n", + " \"GitHub\n", + " View on GitHub\n", + " \n", + "
" + ], + "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": [ + "! 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 +} \ No newline at end of file diff --git a/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/requirements.txt b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/requirements.txt new file mode 100644 index 000000000..983ca6b4c --- /dev/null +++ b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/requirements.txt @@ -0,0 +1 @@ +google-cloud-aiplatform \ No newline at end of file diff --git a/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/Dockerfile b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/Dockerfile new file mode 100644 index 000000000..a269f335b --- /dev/null +++ b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/Dockerfile @@ -0,0 +1,16 @@ +FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-runtime + +RUN apt-get update && \ + apt-get install -y curl gnupg && \ + echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \ + curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \ + apt-get update -y && \ + apt-get install google-cloud-sdk -y + +COPY . /trainer + +WORKDIR /trainer + +RUN pip install -r requirements.txt + +ENTRYPOINT ["python", "-m", "task"] \ No newline at end of file diff --git a/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/__init__.py b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/requirements.txt b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/requirements.txt new file mode 100644 index 000000000..d376b0938 --- /dev/null +++ b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/requirements.txt @@ -0,0 +1,3 @@ +torch==1.8.1 +torchvision==0.9.1 +tensorboard==2.5.0 \ No newline at end of file diff --git a/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/task.py b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/task.py new file mode 100644 index 000000000..8e0f0a5d4 --- /dev/null +++ b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/task.py @@ -0,0 +1,348 @@ +# 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. + +""" +Main program for PyTorch distributed training. +Adapted from: https://github.com/narumiruna/pytorch-distributed-example +""" + +import argparse +import os +import tempfile + +import torch +from torch import distributed +from torch.nn.parallel import DistributedDataParallel +from torch.utils import data +from torch.utils.tensorboard import SummaryWriter + +from torchvision import datasets, transforms + +import utils + +def parse_args(): + + parser = argparse.ArgumentParser() + + # Using environment variables for Cloud Storage directories + # see more details in https://cloud.google.com/vertex-ai/docs/training/code-requirements + parser.add_argument( + '--model-dir', default=os.getenv('AIP_MODEL_DIR'), type=str, + help='a Cloud Storage URI of a directory intended for saving model artifacts') + parser.add_argument( + '--tensorboard-log-dir', default=os.getenv('AIP_TENSORBOARD_LOG_DIR'), type=str, + help='a Cloud Storage URI of a directory intended for saving TensorBoard') + + parser.add_argument( + '--backend', type=str, default='gloo', + help='Use the `nccl` backend for distributed GPU training.' + 'Use the `gloo` backend for distributed CPU training.') + parser.add_argument( + '--init-method', type=str, default='env://', + help='URL specifying how to initialize the package.') + parser.add_argument( + '--world-size', type=int, default=os.environ.get('WORLD_SIZE', 1), + help='The total number of nodes in the cluster. ' + 'This variable has the same value on every node.') + parser.add_argument( + '--rank', type=int, default=os.environ.get('RANK', 0), + help='A unique identifier for each node. ' + 'On the master worker, this is set to 0. ' + 'On each worker, it is set to a different value from 1 to WORLD_SIZE - 1.') + parser.add_argument( + '--epochs', type=int, default=20) + parser.add_argument( + '--no-cuda', action='store_true') + parser.add_argument( + '-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') + + args = parser.parse_args() + + return args + +def distributed_is_initialized(): + if distributed.is_available(): + if distributed.is_initialized(): + return True + return False + +class Average(object): + + def __init__(self): + self.sum = 0 + self.count = 0 + + def __str__(self): + return '{:.6f}'.format(self.average) + + @property + def average(self): + return self.sum / self.count + + def update(self, value, number): + self.sum += value * number + self.count += number + +class Accuracy(object): + + def __init__(self): + self.correct = 0 + self.count = 0 + + def __str__(self): + return '{:.2f}%'.format(self.accuracy * 100) + + @property + def accuracy(self): + return self.correct / self.count + + @torch.no_grad() + def update(self, output, target): + pred = output.argmax(dim=1) + correct = pred.eq(target).sum().item() + + self.correct += correct + self.count += output.size(0) + +class Net(torch.nn.Module): + + def __init__(self, device): + super(Net, self).__init__() + self.fc = torch.nn.Linear(784, 10).to(device) + + def forward(self, x): + return self.fc(x.view(x.size(0), -1)) + +class MNISTDataLoader(data.DataLoader): + + def __init__(self, root, batch_size, train=True): + transform = transforms.Compose([ + transforms.ToTensor(), + transforms.Normalize((0.1307,), (0.3081,)), + ]) + + dataset = datasets.MNIST(root, train=train, transform=transform, download=True) + sampler = None + if train and distributed_is_initialized(): + sampler = data.DistributedSampler(dataset) + + super(MNISTDataLoader, self).__init__( + dataset, + batch_size=batch_size, + shuffle=(sampler is None), + sampler=sampler, + ) + +class Trainer(object): + + def __init__(self, + model, + optimizer, + train_loader, + test_loader, + device, + model_name, + checkpoint_path + ): + self.model = model + self.optimizer = optimizer + self.train_loader = train_loader + self.test_loader = test_loader + self.device = device + self.model_name = model_name + self.checkpoint_path = checkpoint_path + + def save(self, model_dir): + model_path = os.path.join(model_dir, self.model_name) + torch.save(self.model.state_dict(), model_path) + + def fit(self, epochs, is_chief, writer): + + for epoch in range(1, epochs + 1): + + print("Epoch: {}, Training ...".format(epoch)) + train_loss, train_acc = self.train() + + if is_chief: + test_loss, test_acc = self.evaluate() + writer.add_scalar('Loss/train', train_loss.average, epoch) + writer.add_scalar('Loss/test', test_loss.average, epoch) + writer.add_scalar('Accuracy/train', train_acc.accuracy, epoch) + writer.add_scalar('Accuracy/test', test_acc.accuracy, epoch) + torch.save(self.model.state_dict(), self.checkpoint_path) + + print( + 'Epoch: {}/{},'.format(epoch, epochs), + 'train loss: {}, train acc: {},'.format(train_loss, train_acc), + 'test loss: {}, test acc: {}.'.format(test_loss, test_acc), + ) + + def train(self): + + self.model.train() + + train_loss = Average() + train_acc = Accuracy() + + for data, target in self.train_loader: + data = data.to(self.device) + target = target.to(self.device) + + output = self.model(data) + loss = torch.nn.functional.cross_entropy(output, target) + + self.optimizer.zero_grad() + loss.backward() + self.optimizer.step() + + train_loss.update(loss.item(), data.size(0)) + train_acc.update(output, target) + + return train_loss, train_acc + + @torch.no_grad() + def evaluate(self): + self.model.eval() + + test_loss = Average() + test_acc = Accuracy() + + for data, target in self.test_loader: + data = data.to(self.device) + target = target.to(self.device) + + output = self.model(data) + loss = torch.nn.functional.cross_entropy(output, target) + + test_loss.update(loss.item(), data.size(0)) + test_acc.update(output, target) + + return test_loss, test_acc + +def main(): + + args = parse_args() + + local_data_dir = './tmp/data' + local_model_dir = './tmp/model' + local_tensorboard_log_dir = './tmp/logs' + + #TODO: update when gcsfuse ready + gcsfuse_ready = False + model_dir = (gcsfuse_ready and args.model_dir) or local_model_dir + tensorboard_log_dir = (gcsfuse_ready and + args.tensorboard_log_dir) or local_tensorboard_log_dir + + writer = SummaryWriter(tensorboard_log_dir) + + is_chief = args.rank == 0 + + if args.world_size > 1: + print('Initializing distributed backend with {} nodes'.format(args.world_size)) + distributed.init_process_group( + backend=args.backend, + init_method=args.init_method, + 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='') + + if torch.cuda.is_available() and not args.no_cuda: + device = torch.device("cuda:{}".format(args.rank)) + else: + device = torch.device("cpu") + + model = Net(device=device) + if distributed_is_initialized(): + 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) + + # 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)) + else: + model.load_state_dict(torch.load(checkpoint_path)) + + optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate) + + train_loader = MNISTDataLoader( + local_data_dir, args.batch_size, train=True) + test_loader = MNISTDataLoader( + local_data_dir, args.batch_size, train=False) + + trainer = Trainer( + model=model, + optimizer=optimizer, + train_loader=train_loader, + test_loader=test_loader, + device=device, + model_name='mnist.pt', + checkpoint_path=checkpoint_path, + ) + trainer.fit(args.epochs, is_chief, writer) + + if is_chief and model_dir == local_model_dir: + utils.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() + + if is_chief: + os.remove(checkpoint_path) + + if distributed_is_initialized(): + distributed.destroy_process_group() + + return + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/utils.py b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/utils.py new file mode 100644 index 000000000..b44fde47a --- /dev/null +++ b/community-content/pytorch_image_classification_distributed_data_parallel_training_with_vertex_sdk/trainer/utils.py @@ -0,0 +1,31 @@ +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the \"License\"); +# you may not use this file except in compliance with the License.\n", +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an \"AS IS\" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import shutil +import subprocess + +def makedirs(model_dir): + if os.path.exists(model_dir) and os.path.isdir(model_dir): + shutil.rmtree(model_dir) + os.makedirs(model_dir) + return + +def gcs_upload(dir, local_dir, gcs_dir, gcsfuse_ready, local_mode): + if not local_mode and dir == local_dir and not gcsfuse_ready: + subprocess.run(['gsutil', 'cp', '-r', + local_dir, + os.path.dirname(gcs_dir)]) + print(f'{local_dir} is uploaded to {gcs_dir}') + return