文档

提交训练作业

更新时间:

PAI Python SDK提供了更易用的API(即HighLevel API),支持您将训练作业提交到PAI,并使用云上资源运行训练作业。本文为您介绍如何准备训练作业脚本,以及如何使用SDK提交训练作业。

概要介绍

SDK提供了HighLevel的训练API:pai.estimator.Estimator支持您将训练作业提交到PAI,使用Estimator提交训练作业的主要流程包括:

  1. 准备训练作业脚本。

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

  3. 通过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提交训练作业。

准备训练脚本

Estimator支持将您编写的本地训练脚本提交到云上执行,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.Estimatorhyperparameters参数配置了训练作业的超参,超参文件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_TRAINPAI_INPUT_TEST环境变量获取到输入的OSS URI对应的本地路径。

  • 训练作业模型保存:

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

获取PAI提供的公共镜像

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

说明

您可以通过公共镜像列表查看PAI提供的镜像内安装的Python三方库信息。

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"))

安装训练代码的依赖

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

例如在以下示例中,您准备了训练脚本目录train_src,可以在对应的目录下准备所需的 requirements.txt,然后在Estimator构建时指定source_dir="train_src"train_src目录会被打包上传,对应的依赖train_src/requirements.txt会在训练脚本运行前安装到作业执行环境。

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

运行训练作业

您可以通过构建pai.estimator.Estimator指定训练作业的脚本目录、启动脚本、超参和机器资源等,然后通过fit方法提交训练作业。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())

本地执行训练作业

在云上的训练作业调试较为困难,因而Estimator也提供了本地执行的模式,方便您在本地环境中,模拟执行作业,调试相应的脚本。当构建Estimator时,传递的instance_type="local",则对应的训练作业会在本地环境中,通过docker运行对应的作业,模拟训练作业的执行。

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提交的训练作业需要按规范读取超参、获取数据路径,以及写出模型到指定路径。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,分别为modelcheckpoints,分别用于存储模型输出和训练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

  • 本页导读 (1)
文档反馈