<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 PyTorch Image Classification Multi-Node Distributed Data Parallel Training on GPU using Vertex Training with Custom Container
<table align="left">
  <td>
    <a href="https
Step1:  Vertex Training using Vertex SDK and Custom Container
Built Custom Container
Step2:  Initialize Vertex SDK
Step3:  Create a Vertex Tensorboard Instance
Step4:  Option
Step5:  Training Output Artifact
Step6:  Clean Up Artifact
<ASSISTANT_TASK:>
Python Code:
PROJECT_ID = "YOUR PROJECT ID"
BUCKET_NAME = "gs://YOUR BUCKET NAME"
REGION = "YOUR REGION"
SERVICE_ACCOUNT = "YOUR SERVICE ACCOUNT"
! gsutil ls -al $BUCKET_NAME
content_name = "pt-img-cls-multi-node-ddp-cust-cont"
Explanation: PyTorch Image Classification Multi-Node Distributed Data Parallel Training on GPU using Vertex Training with Custom Container
<table align="left">
  <td>
    <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">
      <img src="https://cloud.google.com/ml-engine/images/github-logo-32px.png" alt="GitHub logo">
      View on GitHub
    </a>
  </td>
</table>
Setup
End of explanation
hostname = "gcr.io"
image_name = content_name
tag = "latest"
custom_container_image_uri = f"{hostname}/{PROJECT_ID}/{image_name}:{tag}"
Explanation: Vertex Training using Vertex SDK and Custom Container
Built Custom Container
End of explanation
! pip install -r requirements.txt
from google.cloud import aiplatform
aiplatform.init(
    project=PROJECT_ID,
    staging_bucket=BUCKET_NAME,
    location=REGION,
)
Explanation: Initialize Vertex SDK
End of explanation
content_name = content_name + "-gpu"
tensorboard = aiplatform.Tensorboard.create(
    display_name=content_name,
)
Explanation: Create a Vertex Tensorboard Instance
End of explanation
display_name = content_name
gcs_output_uri_prefix = f"{BUCKET_NAME}/{display_name}"
replica_count = 1
machine_type = "n1-standard-4"
accelerator_count = 4
accelerator_type = "NVIDIA_TESLA_K80"
args = [
    "--backend",
    "nccl",
    "--batch-size",
    "128",
    "--epochs",
    "25",
]
custom_container_training_job = aiplatform.CustomContainerTrainingJob(
    display_name=display_name,
    container_uri=custom_container_image_uri,
)
custom_container_training_job.run(
    args=args,
    base_output_dir=gcs_output_uri_prefix,
    replica_count=replica_count,
    machine_type=machine_type,
    accelerator_count=accelerator_count,
    accelerator_type=accelerator_type,
    tensorboard=tensorboard.resource_name,
    service_account=SERVICE_ACCOUNT,
)
print(f"Custom Training Job Name: {custom_container_training_job.resource_name}")
print(f"GCS Output URI Prefix: {gcs_output_uri_prefix}")
Explanation: Option: Use a Previously Created Vertex Tensorboard Instance
tensorboard_name = "Your Tensorboard Resource Name or Tensorboard ID"
tensorboard = aiplatform.Tensorboard(tensorboard_name=tensorboard_name)
Run a Vertex SDK CustomContainerTrainingJob
End of explanation
! gsutil ls $gcs_output_uri_prefix
Explanation: Training Output Artifact
End of explanation
! gsutil rm -rf $gcs_output_uri_prefix
Explanation: Clean Up Artifact
End of explanation
<END_TASK>