訓練 PyTorch 模型 - Amazon SageMaker AI

本文為英文版的機器翻譯版本,如內容有任何歧義或不一致之處,概以英文版為準。

訓練 PyTorch 模型

本主題會逐步引導您使用 HyperPod 訓練 PyTorch 模型的程序。

在此案例中,讓我們使用 hyp-pytorch-job 範本訓練 PyTorch 模型,透過公開常用參數來簡化任務建立。模型成品會存放在 S3 儲存貯體中,以供日後推論使用。不過,這是選用的,您可以選擇偏好的儲存位置。

建立訓練任務

您可以使用 CLI 或 Python SDK 來訓練模型。

使用 CLI

使用下列命令建立訓練任務:

hyp create hyp-pytorch-job \ --version 1.0 \ --job-name test-pytorch-job \ --image pytorch/pytorch:latest \ --command '["python", "train.py"]' \ --args '["--epochs", "10", "--batch-size", "32"]' \ --environment '{"PYTORCH_CUDA_ALLOC_CONF": "max_split_size_mb:32"}' \ --pull-policy "IfNotPresent" \ --instance-type ml.p4d.24xlarge \ --tasks-per-node 8 \ --label-selector '{"accelerator": "nvidia", "network": "efa"}' \ --deep-health-check-passed-nodes-only true \ --scheduler-type "kueue" \ --queue-name "training-queue" \ --priority "high" \ --max-retry 3 \ --volumes '["data-vol", "model-vol", "checkpoint-vol"]' \ --persistent-volume-claims '["shared-data-pvc", "model-registry-pvc"]' \ --output-s3-uri s3://my-bucket/model-artifacts

重要必要參數說明

  • --job-name:訓練任務的唯一識別符

  • --image:包含訓練環境的 Docker 映像

此命令會啟動名為 的訓練任務test-pytorch-job--output-s3-uri 指定訓練模型成品的存放位置,例如 s3://my-bucket/model-artifacts。請注意此位置,因為您需要它來部署自訂模型。

使用 Python SDK

對於程式設計控制,請使用 SDK。建立 Python 指令碼以啟動相同的訓練任務。

from sagemaker.hyperpod import HyperPodPytorchJob from sagemaker.hyperpod.job import ReplicaSpec, Template, Spec, Container, Resources, RunPolicy, Metadata # Define job specifications nproc_per_node = "1" # Number of processes per node replica_specs = [ ReplicaSpec ( name = "pod", # Replica name template = Template ( spec = Spec ( containers = [ Container ( # Container name name="container-name", # Training image image="448049793756.dkr.ecr.us-west-2.amazonaws.com/ptjob:mnist", # Always pull image image_pull_policy="Always", resources=Resources\ ( # No GPUs requested requests={"nvidia.com/gpu": "0"}, # No GPU limit limits={"nvidia.com/gpu": "0"}, ), # Command to run command=["python", "train.py"], # Script arguments args=["--epochs", "10", "--batch-size", "32"], ) ] ) ), ) ] # Keep pods after completion run_policy = RunPolicy(clean_pod_policy="None") # Create and start the PyTorch job pytorch_job = HyperPodPytorchJob ( # Job name metadata = Metadata(name="demo"), # Processes per node nproc_per_node = nproc_per_node, # Replica specifications replica_specs = replica_specs, # Run policy run_policy = run_policy, # S3 location for artifacts output_s3_uri="s3://my-bucket/model-artifacts" ) # Launch the job pytorch_job.create()

監控您的訓練任務

使用以下命令監控任務的進度:

使用 CLI

# Check job status hyp list hyp-pytorch-job # Get detailed information hyp describe hyp-pytorch-job --job-name test-pytorch-job # View logs hyp get-logs hyp-pytorch-job \ --pod-name test-pytorch-job-pod-0 \ --job-name test-pytorch-job

注意:訓練時間會根據模型複雜性和執行個體類型而有所不同。監控日誌以追蹤進度。

這些命令可協助您驗證任務的狀態並疑難排解問題。一旦任務成功完成,模型成品就會儲存至 s3://my-bucket/model-artifacts

使用 Python SDK

將下列程式碼新增至您的 Python 指令碼:

print("List all pods created for this job:") print(pytorch_job.list_pods()) print("Check the logs from pod0:") print(pytorch_job.get_logs_from_pod(pod_name="demo-pod-0")) print("List all HyperPodPytorchJobs:") print(HyperPodPytorchJob.list()) print("Describe job:") print(HyperPodPytorchJob.get(name="demo").model_dump()) pytorch_job.refresh() print(pytorch_job.status.model_dump())

後續步驟

訓練後,模型成品會存放在您指定的 S3 儲存貯體中 (s3://my-bucket/model-artifacts)。您可以使用這些成品來部署模型。目前,您必須手動管理從訓練到推論的轉換。這包括:

  • 尋找成品:檢查 S3 儲存貯體 (s3://my-bucket/model-artifacts) 以確認訓練過的模型檔案是否存在。

  • 記錄路徑:記下用於推論設定的確切 S3 路徑 (例如 s3://my-bucket/model-artifacts/test-pytorch-job/model.tar.gz)。

  • 部署中的參考:在設定自訂端點時提供此 S3 路徑,以確保載入正確的模型。