人工智能平台 PAI：提交训练作业

计费说明

在提交训练作业时，会调用DLC计算资源执行任务，由此会产生与DLC相关的资源费用。计费详情，请参见DLC计费说明。

概要介绍

SDK提供了HighLevel的训练API pai.estimator.Estimator，支持您将训练作业提交到PAI：

• 使用pai.estimator.Estimator配置使用的训练脚本、作业启动命令、超参、训练镜像以及训练资源等作业信息。

• 通过pai.estimator.Estimator.fit方法，指定训练数据，提交训练作业。

示例代码如下：

from pai.estimator import Estimator

# 通过Estimator配置训练作业。
est = Estimator(
    command="<LaunchCommand>"
    source_dir="<SourceCodeDirectory>"
    image_uri="<TrainingImageUri>"
    instance_type="<TrainingInstanceType>",
    hyperparameters={
        "n_estimators": 500,
        "max_depth": 5,
    },
)

# 指定训练数据，提交训练作业。
est.fit(
    inputs={
        "train_data": "oss://<YourOssBucket>/path/to/train/data/",
    }
)

# 获取输出模型路径。
print(est.model_data())

以下内容为您介绍了提交训练作业的具体操作流程和相关原理介绍。

准备训练脚本和所需依赖

• 准备训练脚本

  Estimator支持将您编写的本地训练脚本提交到PAI上执行，PAI的训练服务会按一定的规范准备训练作业环境，运行训练脚本。训练脚本示例如下：

  import argparse
  import os
  import json
  
  def train(hps, train_data, test_data):
      """用户的模型训练代码"""
      pass
  
  def save_model(model):
      """保存模型"""
      # 通过环境变量获取输出模型路径，默认为/ml/output/model/。
      output_model_path = os.environ.get("PAI_OUTPUT_MODEL")
  
      # 将模型写出到output_model_path路径下。
      pass
  
  def load_hyperparameters():
      """读取输入超参"""
      # 通过环境变量 PAI_CONFIG_DIR 获取超参保存目录，默认为/ml/input/config/。
      hps_path = os.path.join(os.environ.get("PAI_CONFIG_DIR"), "hyperparameters.json")
      with open(hps_path, "r") as f:
          hyperparameters = json.load(f)
      return hyperparameters
  
  def run():
      # 1、加载用户训练作业的超参。
      hps = load_hyperparameters()
      print("Hyperparameters: ", hps)
  
      # 2、训练作业的输入数据。
      # 通过est.fit()方法，用户可以指定存储在NAS或OSS上的训练数据，将数据准备到训练容器中。
      # 训练脚本可以通过环境变量（PAI_INPUT_{CHANNEL_NAME}）获取输入数据的本地路径。
      train_data = os.environ.get("PAI_INPUT_TRAIN")
      test_data = os.environ.get("PAI_INPUT_TEST")
  
      model = train(hps, train_data, test_data)
  
      # 3、作业训练代码，在训练结束之后，写出模型到输出模型路径。
      save_model(model)
  
  
  if __name__ == "__main__":
      run()

  训练脚本需要遵循相应规范，包括读取超参和输入数据，然后在训练结束后写出模型。

  • 训练作业超参：

    当您通过pai.estimator.Estimator的hyperparameters参数配置了训练作业的超参，超参文件hyperparameters.json会被准备到PAI_CONFIG_DIR环境变量指定目录下（默认为/ml/input/config/）。训练作业脚本可以通过读取{PAI_CONFIG_DIR}/hyperparameters.json文件获取相应的超参。

    例如您传递的训练作业超参为hyperparameters={"batch_size": 32, "learning_rate": 0.01}，则超参文件{PAI_CONFIG_DIR}/hyperparameters.json内容为：

    {
      "batch_size": "32",
    	"learning-rate": "0.01"
    }
    

  • 训练作业的输入数据：

    您可以通过pai.estimator.Estimator.fit()的inputs参数指定训练所需的输入数据，输入参数需是一个字典， 其中Key是输入数据的名称（ChannelName），Value是输入数据的存储路径，具体示例如下：

    estimator.fits(
      	inputs={
        		"train": "oss://<YourOssBucket>/train/data/train.csv",
        		"test": "oss://<YourOssBucket>/test/data/",
      	}
    )
    

    对应的输入数据路径会被挂载到/ml/input/data/{ChannelName}路径下，您的训练代码可以通过环境变量PAI_INPUT_{ChannelName}获取到数据的本地路径，您的训练作业脚本可以通过读取本地文件的方式获得指定的输入数据。以上的示例中，我们能通过PAI_INPUT_TRAIN和PAI_INPUT_TEST环境变量获取到输入的OSS URI对应的本地路径。

  • 训练作业模型保存：

    训练作业的输出模型需要保存到指定路径，才能将模型持久化保存。您可以通过PAI_OUTPUT_MODEL环境变量获取到模型保存路径（默认为/ml/output/model），然后将模型保存到对应路径。

• 当您的训练脚本有额外的Python包依赖，训练使用的镜像中没有提供时，您可以通过在训练脚本所在目录中编写requirements.txt，相应的三方库依赖会在您的脚本执行前被安装到作业环境中。

您需要将训练脚本和相关依赖文件存放到指定目录，例如，您在本地创建了目录train_src，您需要将准备好的训练脚本和依赖文件存放到该目录中。后续在Estimator构建时指定source_dir="train_src"，则train_src目录会被打包上传到训练环境的指定路径中。

|-- train_src # 用户指定上传的训练脚本目录。
	|-- requirements.txt # 作业的requirements信息。
	`-- train.py # 训练脚本，用户可以通过 python train.py 的命令拉起脚本
	`-- utils.py

获取PAI提供的公共镜像

您在提交训练作业时，需要指定作业的运行镜像，镜像中需要包含作业脚本的执行依赖，例如使用的机器学习框架，以及相关的第三方依赖。您可以使用保存在阿里云镜像仓库（ACR）中的镜像提交训练作业，也可以使用PAI提供的公共镜像。对于常见的机器学习框架，PAI提供了公共镜像供您使用。您可以使用pai.image.retrieve方法，并指定所需的机器学习框架来获取PAI提供的公共镜像。示例代码如下：

from pai.image import retrieve, list_images, ImageScope


# 获取PAI提供的所有PyTorch训练镜像。
for image_info in list_images(framework_name="PyTorch"):
 print(image_info)

# 获取PAI提供的TensorFlow 2.3版本的CPU训练镜像。
print(retrieve(framework_name="TensorFlow", framework_version="2.3"))

# 获取PAI提供的最新的TensorFlow的GPU训练镜像。
# 通过参数framework_version="latest"，retrieve方法会返回最新的TensorFlow镜像。
print(retrieve(framework_name="TensorFlow", framework_version="latest",
		accelerator_type="GPU"))

# 获取PAI提供的PyTorch 1.12版本的GPU训练镜像。
print(retrieve(framework_name="PyTorch", framework_version="1.12",
 accelerator_type="GPU"))

运行训练作业

将训练作业提交到PAI上运行

您可以通过构建pai.estimator.Estimator指定训练作业的脚本目录、启动脚本、超参和机器资源等，然后通过fit方法提交训练作业。在提交作业之后，会打印训练作业的控制台链接，并持续打印作业的输出日志信息，直到训练作业结束退出（作业状态为成功、失败或被停止）。您可以通过点击作业链接，去控制台查看作业执行详情、日志、机器的资源使用情况，以及训练的Metric等信息。

默认情况下，fit方法在作业执行完成后会退出。您可以使用estimator.model_data()方法获取训练作业输出模型在OSS上的存储路径。

示例代码如下：

from pai.estimator import Estimator
from pai.image import retrieve

# 获取PAI支持的最新PyTorch镜像。
torch_image_uri = retrieve("PyTorch", framework_version="1.12").image_uri

est = Estimator(
    # 训练作业的启动命令。
    command="python train.py",
    # 训练作业脚本，可以是一个本地目录相对路径或是绝对路径，或是OSS上的tar包（例如oss://<YourOssBucket>/your-code-path-to/source.tar.gz）。
    # 当目录中有requirements.txt文件时，对应的依赖会被自动安装，然后再启动用户的训练脚本。
    source_dir="./train_src/",
    # 训练作业使用的镜像。
    image_uri=torch_image_uri,
    # 训练作业使用的机器类型。
    instance_type="ecs.c6.large",
    # 训练作业的超参。
    hyperparameters={
        "n_estimators": 500,
        "objective": "reg:squarederror",
        "max_depth": 5,
    },
    # 训练作业名称前缀，用户提交的训练作业使用的Name为{base_job_name}_{submitted-datetime}。
    base_job_name="example_train_job",
)

# 提交训练作业，同时打印训练作业的Web详情页URL。fit调用默认等待到作业终止（成功、失败、或被停止）。
est.fit()

# 输出的模型路径。
print(est.model_data())

本地执行训练作业

在PAI上提交的训练作业调试较为困难，因此Estimator也提供了本地执行的模式，方便您在本地环境中模拟执行作业，调试相应的脚本。构建Estimator时，若传递参数instance_type="local"，训练作业将在本地通过Docker容器运行，以模拟PAI环境中的作业执行过程。

estimator = Estimator(
    image_uri=image_uri,
    entry_point="train.py",
    # instance_type="local" 表示运行在本地环境。
    instance_type="local",
)

estimator.fit(
    inputs={
        # local模式下，OSS上的数据会被下载到本地，然后挂载到工作容器内。
        "train": "oss://<BucketName>/path-to-data/",
        # local模式下，支持本地文件数据，对应的数据会被挂载到相应的channel目录。
        "test": "/data/to/test/data"
    }
)

# 返回一个本地的模型输出目录。
print(estimator.model_data())

相关文档

• 使用PAI Python SDK训练和部署Pytorch模型的完整操作流程，请参见使用PAI Python SDK训练和部署PyTorch模型。

附录

训练作业预置环境变量

您在PAI提交的训练作业需要按规范读取超参、获取数据路径，以及写出模型到指定路径。PAI的训练服务会将这些信息以环境变量的形式注入到训练作业的容器中，您可以通过环境变量在训练脚本中，或是在Estimator.command指定的训练作业启动命令中，获取超参数、输入数据路径和模型保存路径等信息。

• PAI_HPS_{HYPERPARAMETER_NAME}

  单个训练作业超参的值，会以环境变量的形式注入到训练作业的容器中。对于超参名中，环境变量中不支持的字符（默认的环境变量仅支持使用字母、数字，以及下划线），会被替换为下划线。

  例如您指定了超参{"epochs": 10, "batch-size": 32, "train.learning_rate": 0.001}，对应的环境变量信息如下：

  PAI_HPS_EPOCHS=10
  PAI_HPS_BATCH_SIZE=32
  PAI_HPS_TRAIN_LEARNING_RATE=0.001
  

  您可以在训练启动命令中直接引用这些环境变量，例如：

  est = Estimator(
   command="python train.py --epochs $PAI_HPS_EPOCHS --batch-size $PAI_HPS_BATCH_SZIE",
   hyperparameters={
   "epochs": 10,
   "batch-size": 32,
   },
   # more arguments for estimator..
  )
  

  以上的方式传递的参数，训练脚本train.py可以通过标准库argparse获取输入参数。

• PAI_USER_ARGS

  训练作业的所有超参信息，会以PAI_USER_ARGS环境变量，使用--{hyperparameter_name} {hyperparameter_value}的形式，注入到训练作业的容器中。

  例如训练作业指定了超参hyperparameters={"epochs": 10, "batch-size": 32, "learning-rate": 0.001}，则PAI_USER_ARGS环境变量的值为：

  PAI_USER_ARGS="--epochs 10 --batch-size 32 --learning-rate 0.001"
  

  您可以在启动命令中引用环境变量，例如以下的实例中，训练作业脚本会以python train.py --epochs 10 --batch-size 32 --learning-rate 0.001的命令执行。

  est = Estimator(
      command="python train.py $PAI_USER_ARGS",
      hyperparameters={
          "epochs": 10,
          "learning-rate": 0.001
          "batch-size": 32,
      },
      # more arguments for estimator..
  )

• PAI_HPS

  您的训练作业的超参信息，会以JSON格式，通过PAI_HPS环境变量注入到训练作业的容器中。

  例如您传递了超参{"epochs": 10, "batch-size": 32}，则PAI_HPS环境变量的值为：

  PAI_HPS={"epochs": 10, "batch-size": 32}
  

• PAI_INPUT_{channel_name}

  训练作业的输入数据，会以挂载的形式挂载到训练作业上，您可以通过读取本地文件的方式读取到OSS和NAS上的数据。对于每一个输入的数据，会以PAI_INPUT_{channel_name}的环境变量注入到训练作业的容器中。

  例如您的执行代码为est.fit(inputs={"train": "oss://<YourOssBucket>/path-to-data/", "test": "oss://<YourOssBucket>/path-to/data/test.csv"})，对应的环境变量如下：

  PAI_INPUT_TRAIN=/ml/input/data/train/
  PAI_INPUT_TEST=/ml/input/data/test/test.csv
  

  对应的数据存储路径会被挂载到容器中，您可以通过这些本地路径信息，直接读取到输入的数据。

  说明

  PAI_INPUT_{ChannelName}指向您传入的数据路径，如果您指定了一个OSS目录（以 / 结尾），则PAI的训练服务会将输入存储作为目录进行挂载，环境变量指向对应的数据目录。如果您传递了一个OSS文件路径，PAI的训练服务会挂载对应的文件目录，环境变量会指OSS URI对应的实际文件。

• PAI_OUTPUT_{channel_name}

  默认训练作业会创建两个输出Channel，分别为model和checkpoints，分别用于存储模型输出和训练checkpoints。每一个Channel对应一个OSS URI，以及对应的挂载路径。您可以通过PAI_OUTPUT_{channel_name}环境变量，获取到对应的文件路径。

  PAI_OUTPUT_MODEL=/ml/output/model/
  PAI_OUTPUT_CHECKPOINTS=/ml/output/checkpoints/
  

  通过将需要保存的模型或checkpoints，保存到这些路径下，PAI的训练服务会自动将这些文件上传到对应的OSS路径下。

训练作业目录结构

完整的训练作业的输入输出目录结构示例如下：

/ml
|-- usercode                        # 用户代码加载到/ml/usercode目录，这里也是用户代码的工作目录。可以通过环境变量 `PAI_WORKING_DIR` 获得.
|   |-- requirements.txt
|   `-- train.py
|-- input                           # 作业输入数据和配置信息。
|   `-- config                      # config目录包含了作业的配置信息，可以通过PAI_CONFIG_DIR获取。
|       |-- hyperparameters.json    # 训练作业超参。
|   `-- data                        # 作业的InputChannels：以下目录包含了两个channel: train_data和test_data。
|       |-- test_data
|       |   `-- test.csv
|       `-- train_data
|           `-- train.csv
`-- output                          # 作业的输出Channels：默认包含两个OutputChannel: model/checkpoints。
        `-- model                   # 通过环境变量PAI_OUTPUT_{CHANNEL_NAME}可以获输出路径。
        `-- checkpoints