mirror of
https://github.com/GoogleCloudPlatform/vertex-ai-samples.git
synced 2026-09-26 14:42:04 +00:00
content: add pytorch-img-cls-multi-node-ddp-cust-cont (#37)
This commit is contained in:
Morgan Du
+554
@@ -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": [
|
||||
"<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_gloo_vertex_training_with_custom_container.ipynb\">\n",
|
||||
" <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\">\n",
|
||||
" View on GitHub\n",
|
||||
" </a>\n",
|
||||
" </td>\n",
|
||||
"</table>"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"source": [
|
||||
"## Setup"
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"PROJECT_ID = \"YOUR PROJECT ID\"\n",
|
||||
"BUCKET_NAME = \"gs://YOUR BUCKET NAME\"\n",
|
||||
"REGION = \"YOUR REGION\"\n",
|
||||
"SERVICE_ACCOUNT = \"YOUR SERVICE ACCOUNT\""
|
||||
],
|
||||
"metadata": {
|
||||
"collapsed": false,
|
||||
"pycharm": {
|
||||
"name": "#%%\n"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"! 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
|
||||
}
|
||||
+391
@@ -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": [
|
||||
"<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"
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"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
|
||||
}
|
||||
+1
@@ -0,0 +1 @@
|
||||
google-cloud-aiplatform
|
||||
+16
@@ -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"]
|
||||
+3
@@ -0,0 +1,3 @@
|
||||
torch==1.8.1
|
||||
torchvision==0.9.1
|
||||
tensorboard==2.5.0
|
||||
+348
@@ -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()
|
||||
+31
@@ -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
|
||||
Reference in New Issue
Block a user