容器服务运行 PyTorch 训练任务-云原生服务指南-文档中心-腾讯云

本文为您介绍如何运行 PyTorch 训练任务。
前提条件
AI 环境中已安装 PyTorch Operator。
AI 环境中有 GPU 资源。
操作步骤
以下操作指南参考 PyTorch-Operator 官方提供的分布式训练 案例。
准备训练代码
本示例使用 Kubeflow 官方提供的示例代码 mnist.py。
制作训练镜像
制作训练镜像的过程较为简单。您只需要基于 PyTorch 1.0 的官方镜像，将代码复制到镜像内，并配置好entrypoint即可。如果您没有配置 entrypoint，也可以在提交 PyTorchJob 时配置启动命令。
注意
训练代码基于 PyTorch 1.0 构建，由于 PyTorch 在不同版本之间可能存在 API 不兼容的问题，因此在其他版本的 PyTorch 环境下，您可能需要自行调整代码。
任务提交
1. 准备一个 PyTorchJob 的 YAML 文件，定义一个 Master Worker 和一个 Worker。
注意：
您需要将上传后的训练镜像地址替换  <训练镜像> 所在占位。
由于在资源配置中设置了 GPU 资源，请在args为训练配置的backend为"nccl"；在未使用 （Nvidia）GPU 的任务中，请使用其他 backend（例如 gloo）。
apiVersion: "kubeflow.org/v1"
kind: "PyTorchJob"
metadata:
  name: "pytorch-dist-mnist-nccl"
spec:
  pytorchReplicaSpecs:
    Master:
      replicas: 1
      restartPolicy: OnFailure
      template:
        metadata:
          annotations:
            sidecar.istio.io/inject: "false"
        spec:
          containers:
            - name: pytorch
              image: <训练镜像>
              args: ["--backend", "nccl"]
              resources: 
                limits:
                  nvidia.com/gpu: 1
    Worker:
      replicas: 1
      restartPolicy: OnFailure
      template:
        metadata:
          annotations:
            sidecar.istio.io/inject: "false"
        spec:
          containers: 
            - name: pytorch
              image: <训练镜像>
              args: ["--backend", "nccl"]
              resources: 
                limits:
                  nvidia.com/gpu: 1
2. 执行以下命令，通过 kubectl 提交 PyTorchJob：
kubectl create -f ./pytorch_job_mnist_nccl.yaml
3. 执行以下命令，查看该 PyTorchJob：
kubectl get -o yaml pytorchjobs pytorch-dist-mnist-nccl
4. 执行以下命令，查看 PyTorch 任务创建的相关 Pod：
kubectl get pods -l pytorch_job_name=pytorch-dist-mnist-nccl  
﻿
运行 PyTorch 训练任务

本页目录：

前提条件

操作步骤

准备训练代码

制作训练镜像

任务提交
