如何使用 SageMaker 训练人机语音交互模型

发布于: Oct 10, 2022

在完成有关实验环境的准备工作后，接下来，我们就进入正题，利用SageMaker 完成人机语音交互模型训练工作。
WeNet 支持训练多种模型，如 Conformer、Transformer 等，这里我们会以unified transformer 为例展示整个训练流程。对于训练数据，WeNet 同样支持多种来源，只需要在训练的时候，按照格式整理数据即可，如 AIShell-1、AIShell-2 及 LibriSpeech 等，这里，我们会以 AIShell-1 为例。

数据下载

我们首先需要将训练数据下载到本地 FSx 存储中，在 notebook 中执行命令：

cd /fsx/wenet/examples/aishell/s0 && \bash run.sh --stage -1 --stop_stage -1 --data /fsx/asr-data/OpenSLR/33

数据会自动被下载到 /fsx/asr-data/OpenSLR/33 目录中，下载完成后的状态为：

sh-4.2$ ls /fsx/asr-data/OpenSLR/33
data_aishell  data_aishell.tgz  resource_aishell  resource_aishell.tgz

数据预处理

接下来，我们需要将数据整理为 WeNet 所需的格式。这里我们借助 SageMaker 来执行数据预处理的逻辑。

挂载 FSx 文件系统到数据预处理容器

在前面我们提到，模型训练所需的数据已经存放在 FSx 文件系统中，我们在通过 SageMaker 处理数据的时候，需要把此 FSx 文件系统挂载到容器里。挂载文件系统的代码如下：

from sagemaker.inputs import FileSystemInputfrom sagemaker.pytorch.estimator import PyTorch

file_system_id = 'fs-0f8a3xxxxf47b6ff8'
file_system_path = '/yobzhbmv'
file_system_access_mode = 'rw'
file_system_type = 'FSxLustre'

security_group_ids = ['sg-04acfcxxxx929ee4e']
subnets= ['subnet-07ce0abxxxxcfeb25']

file_system_input_train = FileSystemInput(file_system_id=file_system_id,
                                  file_system_type=file_system_type,
                                  directory_path=file_system_path,
                                  file_system_access_mode=file_system_access_mode)

需要注意，subnets 参数中指定的子网，需要有访问 S3 等服务的能力，您可以选择使用私有子网，并为子网指定到 NAT 网关的默认路由。security_group_ids 指定的安全组会被绑定到 SageMaker 启动的实例上，需要有访问FSx服务的能力。

启动数据预处理作业

至此，我们通过指定文件系统的 id、文件系统的路径、读写模式等信息，定义好了需要挂载的文件系统。接下来，就可以设置数据处理的时候，运行环境及需要传递的参数信息。代码如下：

hp= {
    'stage': 0, 'stop_stage': 3, 'train_set':'train', 
    'trail_dir':'/opt/ml/input/data/train/sm-train/trail0', 
    'data': '/opt/ml/input/data/train/asr-data/OpenSLR/33',
    'shared_dir': '/opt/ml/input/data/train/shared'}

estimator=PyTorch(
    entry_point='examples/aishell/s0/sm-run.sh',
    image_uri=training_repository_uri,
    instance_type='ml.c5.xlarge',
    instance_count=1,
    source_dir='.',
    role=role,
    hyperparameters=hp,
    
    subnets=subnets,
    security_group_ids=security_group_ids,
    
    debugger_hook_config=False,
    disable_profiler=True)

我们通过 image_uri 参数指定数据处理代码运行的容器环境，instance_type 指定需要的实例类型，instance_count 指定需要的实例数量，hyperparameters 指定需要传递的超参数。
接下来，就可以通过一行命令启动指定的计算资源，并执行数据处理逻辑。

estimator.fit(inputs={'train': file_system_input_train})

我们通过 inputs 参数设置了容器运行时的数据输入信息，SageMaker 支持多种数据来源，如本地文件(file://)，S3路径(s3://bucket/path)及文件系统(FSx 或者EFS)。这里，我们的 FSx 文件系统会被映射到容器的 /opt/ml/input/data/train 目录下，train 为自定义的channel名称，其他常见的 channel 包括 test，validation 等。SageMaker 中具体的路径映射规则可以参考[1]。

查看处理后的数据

处理完成之后，会在 trail_dir 及 shared_dir 目录下创建对应的文件。在 Notebook 实例上执行命令，具体如下：

tree -L 3 /fsx/sm-train/trail0	tree -L 3 /fsx/sm-train/shared

启动模型训练作业

至此，我们已经准备好了训练数据。接下来，我们就可以进入模型训练阶段了。我们会展示本地训练和全托管实例训练两种训练模式。

本地训练模式

在模型研发过程中，算法人员需要反复调整代码逻辑，如果每次代码调整就打包一个 docker 镜像就显得很麻烦，因此，您可以先通过SageMaker的本地训练模式，来调试代码。本地训练模式会直接在 Notebook 所在实例中启动对应的容器并执行训练逻辑，并自动将数据映射给容器。有关本地模式训练的细节，可以参考文档[3]，这里我们使用的本地训练代码如下：

instance_type='local_gpu'
instance_count = 1
CUDA_VISIBLE_DEVICES='0'

hp= {
    'stage': 4, 'stop_stage': 4, 'train_set':'train', 
    'data': data_dir, 'trail_dir': trail_dir, 'shared_dir': shared_dir,
    'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 
    'num_nodes': instance_count}

estimator=PyTorch( 
    entry_point='examples/aishell/s0/sm-run.sh',
    image_uri=training_repository_uri,
    instance_type =instance_type,
    instance_count=instance_count,
    source_dir='.',
    role=role,
    hyperparameters=hp,
    
    subnets=subnets,
    security_group_ids=security_group_ids,
    
    debugger_hook_config=False,
    disable_profiler=True)


estimator.fit({'train': 'file:///fsx'})

代码的输出如下：

Creating 2n0im72bz3-algo-1-tpyyu ... 
Creating 2n0im72bz3-algo-1-tpyyu ... done
Attaching to 2n0im72bz3-algo-1-tpyyu
…
2n0im72bz3-algo-1-tpyyu | Invoking script with the following command:
2n0im72bz3-algo-1-tpyyu | 
2n0im72bz3-algo-1-tpyyu | /bin/sh -c ./examples/aishell/s0/sm-run.sh --CUDA_VISIBLE_DEVICES 0 --data /opt/ml/input/data/train/asr-data/OpenSLR/33 --num_nodes 1 --shared_dir /opt/ml/input/data/train/sm-train/shared --stage 4 --stop_stage 4 --trail_dir /opt/ml/input/data/train/sm-train/trail0 --train_set train
…
2n0im72bz3-algo-1-tpyyu | algo-1-tpyyu: 2021-06-24 15:50:09,408 INFO     [checkpoint.py:33] Checkpoint: save to checkpoint /opt/ml/input/data/train/sm-train/trail0/exp/unified_transformer/init.pt
2n0im72bz3-algo-1-tpyyu | algo-1-tpyyu: 2021-06-24 15:50:09,669 INFO     [train.py:228] Epoch 0 TRAIN info lr 8e-08
2n0im72bz3-algo-1-tpyyu | algo-1-tpyyu: 2021-06-24 15:50:09,670 INFO     [executor.py:32] using accumulate grad, new batch size is 1 timeslarger than before
2n0im72bz3-algo-1-tpyyu | algo-1-tpyyu: 2021-06-24 15:50:12,560 DEBUG    [executor.py:103] TRAIN Batch 0/7507 loss 417.150146 loss_att 148.725983 loss_ctc 1043.473145 lr 0.00000008 rank 0

上述参数中，source_dir 指定的路径会被打包上传到 S3，然后，下载到容器实例中。这样的话，我们每次的代码变更都可以直接体现在容器中。
此外，在使用本地训练模式时，SageMaker 会借助本地的 docker-compose 启动对应的训练任务，您可以在 /tmp 目录下找到相关的 docker-compose 文件，如 /tmp/tmp6y009akq，我们可以观察到如下内容：

sh-4.2$ tree /tmp/tmp6y009akq
/tmp/tmp6y009akq
├── artifacts
├── docker-compose.yaml
├── model
└── output
    └── data其中，docker-compose.yaml包含了相关的配置信息，内容如下：

sh-4.2$ cat /tmp/tmp6y009akq/docker-compose.yaml 
networks:
  sagemaker-local:
    name: sagemaker-local
services:
  algo-1-tpyyu:
    command: train
    container_name: 2n0im72bz3-algo-1-tpyyu
    environment:
    - AWS_REGION=us-east-1
    - TRAINING_JOB_NAME=sagemaker-wenet-2021-06-24-15-49-58-018
    image: <your-aws-account-id>.dkr.ecr.us-east-1.amazonaws.com/sagemaker-wenet:training-pip-pt181-py38
    networks:
      sagemaker-local:
        aliases:
        - algo-1-tpyyu
    stdin_open: true
    tty: true
    volumes:
    - /tmp/tmp6y009akq/algo-1-tpyyu/output:/opt/ml/output
    - /tmp/tmp6y009akq/algo-1-tpyyu/output/data:/opt/ml/output/data
    - /tmp/tmp6y009akq/algo-1-tpyyu/input:/opt/ml/input
    - /tmp/tmp6y009akq/model:/opt/ml/model
    - /opt/ml/metadata:/opt/ml/metadata
    - /fsx:/opt/ml/input/data/train
version: '2.3'

可以看到，docker-compose 通过 volumes 参数，将本地的路径映射为容器里的目录，而不需要执行训练数据的二次复制。

托管训练模式

在确定代码逻辑无误后，我们可以很容易通过修改参数的方式，使用托管的实例开启真正的训练任务。
这里，我们只需要调整实例类型、需要的实例数量及数据输入方式。我们以 2 台 ml.p3.8xlarge 的实例为例，其各自包含 4 张 Tesla V100 显卡，共8张显卡。
训练代码如下：

instance_type='ml.p3.8xlarge'
instance_count = 2
CUDA_VISIBLE_DEVICES='0,1,2,3'

hp= {
    'stage': 4, 'stop_stage': 4, 'train_set':'train', 
    'data': data_dir, 'trail_dir': trail_dir, 'shared_dir': shared_dir,
    'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 
    'ddp_init_protocol': 'tcp',
    'num_nodes': instance_count}

estimator=PyTorch( 
    entry_point='examples/aishell/s0/sm-run.sh',
    image_uri=training_repository_uri,
    instance_type =instance_type,
    instance_count=instance_count,
    source_dir='.',
    role=role,
    hyperparameters=hp,
    
    subnets=subnets,
    security_group_ids=security_group_ids,
    
    debugger_hook_config=False,
    disable_profiler=True,
    environment={
        'NCCL_SOCKET_IFNAME': 'eth0',
        'NCCL_IB_DISABLE': 1
    })

estimator.fit(inputs={'train': file_system_input_train})

其中，参数 CUDA_VISIBLE_DEVICES 需设定为训练实例的 GPU 卡数量。如果仅有一张 GPU 显卡，则其值为 ’0’。
这里需要注意的是，撰写本文时，SageMaker 训练任务在挂载FSx时，还不支持指定挂载选项 flock，导致无法使用基于file的分布式初始化方法。因此，我们简单调整 WeNet 的训练代码，转而使用基于 TCP 的初始化方法，来继续模型训练。
您还可以观察到，我们传入了 environment 参数，其表示设定容器中对应的环境变量。由于 SageMaker 拉起的训练实例会包含不止一个网卡，因此，我们需要通过 NCCL_SOCKET_IFNAME 环境变量，将 NCCL 使用的网卡设定为 eth0。
此外，SageMaker 支持使用竞价实例来训练模型，以有效降低成本，您可以参考文档[4]查看使用方法。

模型文件

在训练完成之后，会在您设定的目录生成对应的模型文件，本文为 /fsx/sm-train/trail0/exp/unified_transformer 目录。

如果您需要导出支持序列化和优化的(TorchScript)模型，则可以调整hp变量中的 stage及 stop_stage，通过本地模式执行训练代码即可。有关TorchScript，可以参考[5]。
相关代码逻辑如下：

instance_type='local_gpu'
…
hp= {
    'stage': 5, 'stop_stage': 6, 'train_set':'train', 
…}

estimator=PyTorch(
…)

estimator.fit({'train':'file:///fsx'})

执行完成之后，会在上述目录生成对应的模型文件 final.zip 及量化模型 final_quant.zip 文件。
现在，我们已经完成了一次模型训练工作。我们知道，想要得到一个满足当下需求的模型，需要经历多次试验，多次迭代及训练。您可以通过上述方法，在 SageMaker 上快速尝试不同的超参数或者其他的算法，而无需考虑如何配置机器学习的基础环境等运维相关工作。