发布于: 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 文件系统中,我们在通过 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 上快速尝试不同的超参数或者其他的算法,而无需考虑如何配置机器学习的基础环境等运维相关工作。

至此,我们已经得到了训练好的模型文件。您可以通过 SageMaker 部署模型,也可以通过其他方式部署。在后续的文章中,我们会详细介绍如何在 Amazon Web Services 部署训练好的模型。

相关文章