Amazon SageMaker新玩法——定製你的語音識別模型

前言

通過語音方式與機器進行互動可以在很多場景下提高效率，也是當下人工智慧領域內研究的熱點之一。語音識別技術的應用場景可以劃分為以車載語音助手為例的車載場景、以智慧家居裝置的家庭場景等。要實現人與機器間的語音互動，需要首先讓機器能夠識別聲音內容，但通用的語音識別服務無法完全滿足不同場景下的需求，因此客戶需要根據自己的需求訓練模型。

本文會為大家展示如何使用Amazon SageMaker服務訓練自己的語音識別模型，我們選擇了一個開源的語音識別專案WeNet作為示例。

Amazon SageMaker是一項完全託管的機器學習服務，涵蓋了資料標記、資料處理、模型訓練、超參調優、模型部署及持續模型監控等基本流程；也提供自動打標籤，自動機器學習，監控模型訓練等高階功能。其通過全託管的機器學習基礎設施和對主流框架的支援，可以降低客戶機器學習的整體擁有成本。

WeNet是一個面向工業級產品的開源端到端語音識別解決方案，同時支援流式及非流式識別，並能高效執行於雲端及嵌入式端。模型在訓練的過程中，需要用到大量的計算資源，我們可以藉助Amazon SageMaker非常方便的啟動包含多臺完全託管的訓練例項叢集，加速訓練過程。

? 想要了解更多亞馬遜雲科技最新技術釋出和實踐創新，敬請關注在上海、北京、深圳三地舉辦的2021亞馬遜雲科技中國峰會！點選圖片報名吧～

準備工作

在開始訓練模型之前，我們需要做一些準備工作，包括準備FSx檔案系統以存放訓練過程中的資料、建立Amazon SageMaker Notebook作為實驗環境、在筆記本中掛載FSx檔案系統、準備實驗程式碼，以及準備資料處理及模型訓練的執行環境(Docker映象)並把映象推送到Amazon ECR(Elastic Container Registry)中。

本文中的實驗內容均使用us-east-1區域中的服務完成，您可以自行使用其他區域。

建立FSx for Lustre儲存

以往，在Amazon SageMaker中訓練模型一般使用Amazon Simple Storage Service（Amazon S3)作為儲存，現在，Amazon SageMaker做模型訓練時已經支援多種資料來源，比如Amazon FSx for Lustre和Amazon Elastic File System (EFS)。Amazon SageMaker通過直接讀取儲存在EFS或者FSx for Luster上的資料來加快訓練模型時資料載入進度。

FSx for Lustre支援從Amazon S3中匯入資料，以及將資料匯出到Amazon S3，如果您的資料已經存放在Amazon S3中，FSx for Lustre以透明方式將物件顯示為檔案。同一個FSx檔案系統還可用於多個Amazon SageMaker訓練任務，省去多次重複下載訓練資料集的時間。

這裡，我們會選擇使用FSx for Lustre作為主要的資料儲存。接下來，我們會建立一個FSx for Lustre儲存。

建立基於Amazon S3的FSx for Lustre

在“網路和安全性”處設定VPC，子網組和安全組，並確認安全組入站規則是否允許了埠998的流量。

在“資料儲存庫匯入/匯出”處選擇“ 從Amazon S3匯入資料及將資料匯出到Amazon S3”並指定Amazon S3訓練資料所在的儲存桶和路徑。

建立完成後，點選“掛載”按鈕就會彈出掛載此檔案系統的步驟，稍後我們會在Amazon SageMaker Notebook中使用到。

建立Amazon SageMaker Notebook

選擇筆記本例項型別，這裡我們選擇一臺ml.p3.8xlarge的機器，其包含4張Tesla V100 GPU卡。您可以選擇其他GPU機器，如果您不需要GPU卡，也可以選擇CPU機器。

此外，您可以自行決定筆記本例項的卷大小，如本例項選擇了100GB的儲存。您可以在後續調整此儲存的大小。

選擇新建IAM角色，包含所需的許可權，如下圖：

網路部分，選擇FSx所在VPC以及公有子網即可，安全組需要允許Amazon SageMaker訪問FSx。

在筆記本中掛載FSx儲存

在筆記本控制檯頁面，點選“開啟JupyterLab”。

在Launcher頁面，點選“Terminal”以建立一個新的命令列終端。根據“建立基於Amazon S3的FSx”章節中提示的步驟，在命令終端中安裝Lustre客戶端，並執行掛載命令。

此外，您還可以配置筆記本生命週期策略，在建立或者啟動Notebook例項的時候，實現筆記本自動掛載FSx檔案系統，參考文件[2]。

下載WeNet原始碼

在上一步中的命令列終端，執行如下命令，將完成程式碼下載。

1sudo chown ec2-user.ec2-user /fsx
2
3ln -s /fsx /home/ec2-user/SageMaker/fsx
4
5cd ~/SageMaker/fsx
6
7git clone -b sagemaker https://github.com/chen188/wenet

這裡，我們建議您將試驗相關檔案都放置在~/Amazon SageMaker目錄下，該目錄下的資料在Notebook例項關機之後依然可以單獨存在。

您可以開啟Notebook檔案
/fsx/wenet/examples/aishell/s0/SM-WeNet.ipynb，
後續的命令您都可以在此筆記本中找到。

準備Docker 映象

在Amazon SageMaker中，很多工都是基於Docker 映象實現，如資料預處理、模型訓練及模型託管等。採用Docker映象可以極大程度保證環境的一致性，並且降低環境預置的運維成本。

接下來，我們會需要構建自己的Docker映象，來實現資料格式轉換、模型訓練。Amazon Web Service已經提供了一些通用的Deep Learning Container(DLC)環境，具體列表可以參考[6]。但是其中尚未包含TorchAudio包，此時，我們可以選擇基於開源版本構建執行環境。

該映象基於Ubuntu來構建，並安裝pytorch 1.8.1、torchaudio及其他相關依賴。

檔案/fsx/wenet/Dockerfile：

1FROM ubuntu:latest
2ENV DEBIAN_FRONTEND=noninteractive
3ENV PATH /opt/conda/bin:$PATH
4
5RUN apt-get update --fix-missing && \
6    apt-get install -y gcc net-tools && \
7    apt-get install -y --no-install-recommends wget bzip2 ca-certificates libglib2.0-0 libxext6 libsm6 libxrender1 git mercurial subversion && \
8    apt-get clean && \
9    rm -rf /var/lib/apt/lists/* && \
10    wget --quiet https://repo.anaconda.com/archive/Anaconda3-2021.05-Linux-x86_64.sh -O ~/anaconda.sh && \
11    /bin/bash ~/anaconda.sh -b -p /opt/conda && \
12    rm ~/anaconda.sh && \
13    ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
14    echo ". /opt/conda/etc/profile.d/conda.sh" >> ~/.bashrc && \
15    echo "conda activate base" >> ~/.bashrc && \
16    find /opt/conda/ -follow -type f -name '*.a' -delete && \
17    find /opt/conda/ -follow -type f -name '*.js.map' -delete && \
18    /opt/conda/bin/conda clean -afy
19
20COPY ./requirements.txt /tmp/
21
22RUN pip install -r /tmp/requirements.txt && \
23    pip install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html && \
24    pip install sagemaker-training && \
25rm /tmp/requirements.txt

您可以注意到，我們額外安裝了Amazon SageMaker-training包，來提供映象對Amazon SageMaker訓練功能的支援。

構建映象並推送到ECR

ECR是Amazon完全託管的容器登錄檔服務，我們可以將構建好的映象推送到ECR，後續Amazon SageMaker在訓練或者託管模型的時候，會從這裡下載對應的映象。

1import boto3
2account_id = boto3.client('sts').get_caller_identity().get('Account')
3region = boto3.Session().region_name
4ecr_repository = 'sagemaker-wenet'
5
6# 登入ECR服務
7!aws ecr get-login-password --region {region} | docker login --username AWS --password-stdin {account_id}.dkr.ecr.{region}.amazonaws.com
8
9# 訓練映象
10training_docker_file_path = '/fsx/wenet'
11!cat $training_docker_file_path/Dockerfile
12
13tag = ':training-pip-pt181-py38'
14training_repository_uri = '{}.dkr.ecr.{}.amazonaws.com/{}'.format(account_id, region, ecr_repository + tag)
15print('training_repository_uri: ', training_repository_uri)
16
17!cd $training_docker_file_path && docker build -t "$ecr_repository$tag" .
18!docker tag {ecr_repository + tag} $training_repository_uri
19!docker push $training_repository_uri
20Python

使用Amazon SageMaker訓練模型

現在，我們已經完成實驗環境準備工作，接下來，我們就進入正題，利用Amazon SageMaker完成模型訓練工作。

WeNet支援訓練多種模型，如Conformer、Transformer等，這裡我們會以unified transformer為例展示整個訓練流程。對於訓練資料，WeNet同樣支援多種來源，只需要在訓練的時候，按照格式整理資料即可，如AIShell-1、AIShell-2及LibriSpeech等，這裡，我們會以AIShell-1為例。

格式
https://wenet-e2e.github.io/w...

資料下載

我們首先需要將訓練資料下載到本地FSx儲存中，在notebook中執行命令：

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

資料會自動被下載到/fsx/asr-data/OpenSLR/33目錄中，下載完成後的狀態為：

1sh-4.2$ ls /fsx/asr-data/OpenSLR/33
2data_aishell  data_aishell.tgz  resource_aishell  resource_aishell.tgz
3Bash

資料預處理

接下來，我們需要將資料整理為WeNet所需的格式。這裡我們藉助Amazon SageMaker來執行資料預處理的邏輯。

掛載FSx檔案系統到資料預處理容器

在前面我們提到，模型訓練所需的資料已經存放在FSx檔案系統中，我們在通過Amazon SageMaker處理資料的時候，需要把此FSx檔案系統掛載到容器裡。掛載檔案系統的程式碼如下：

1from sagemaker.inputs import FileSystemInput
2from sagemaker.pytorch.estimator import PyTorch
3
4file_system_id = 'fs-0f8a3xxxxf47b6ff8'
5file_system_path = '/yobzhbmv'
6file_system_access_mode = 'rw'
7file_system_type = 'FSxLustre'
8
9security_group_ids = ['sg-04acfcxxxx929ee4e']
10subnets= ['subnet-07ce0abxxxxcfeb25']
11
12file_system_input_train = FileSystemInput(file_system_id=file_system_id,
13                                  file_system_type=file_system_type,
14                                  directory_path=file_system_path,
15                                  file_system_access_mode=file_system_access_mode)
16Python

需要注意，subnets引數中指定的子網，需要有訪問Amazon S3等服務的能力，您可以選擇使用私有子網，併為子網指定到NAT閘道器的預設路由。

security_group_ids指定的安全組會被繫結到Amazon SageMaker啟動的例項上，需要有訪問FSx服務的能力。

啟動資料預處理作業

至此，我們通過指定檔案系統的id、檔案系統的路徑、讀寫模式等資訊，定義好了需要掛載的檔案系統。接下來，就可以設定資料處理的時候，執行環境及需要傳遞的引數資訊。程式碼如下：

1hp= {
2    'stage': 0, 'stop_stage': 3, 'train_set':'train', 
3    'trail_dir':'/opt/ml/input/data/train/sm-train/trail0', 
4    'data': '/opt/ml/input/data/train/asr-data/OpenSLR/33',
5    'shared_dir': '/opt/ml/input/data/train/shared'
6}
7
8estimator=PyTorch(
9    entry_point='examples/aishell/s0/sm-run.sh',
10    image_uri=training_repository_uri,
11    instance_type='ml.c5.xlarge',
12    instance_count=1,
13    source_dir='.',
14    role=role,
15    hyperparameters=hp,
16
17    subnets=subnets,
18    security_group_ids=security_group_ids,
19
20    debugger_hook_config=False,
21    disable_profiler=True
22)
23Python

我們通過image_uri引數指定資料處理程式碼執行的容器環境，instance_type指定需要的例項型別，instance_count指定需要的例項數量，hyperparameters指定需要傳遞的超引數。

接下來，就可以通過一行命令啟動指定的計算資源，並執行資料處理邏輯。

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

我們通過inputs引數設定了容器執行時的資料輸入資訊，Amazon SageMaker支援多種資料來源，如本地檔案(file://)，Amazon S3路徑(s3://bucket/path)及檔案系統(FSx或者EFS)。這裡，我們的FSx檔案系統會被對映到容器的 /opt/ml/input/data/train 目錄下，train為自定義的channel名稱，其他常見的channel包括test，validation等。Amazon SageMaker中具體的路徑對映規則可以參考[1]。

檢視處理後的資料

處理完成之後，會在trail_dir及shared_dir目錄下建立對應的檔案。在Notebook例項上執行命令，具體如下：

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

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

啟動模型訓練作業

至此，我們已經準備好了訓練資料。接下來，我們就可以進入模型訓練階段了。我們會展示本地訓練和全託管例項訓練兩種訓練模式。

本地訓練模式

在模型研發過程中，演算法人員需要反覆調整程式碼邏輯，如果每次程式碼調整就打包一個docker映象就顯得很麻煩，因此，您可以先通過Amazon SageMaker的本地訓練模式，來除錯程式碼。本地訓練模式會直接在Notebook所在例項中啟動對應的容器並執行訓練邏輯，並自動將資料對映給容器。有關本地模式訓練的細節，可以參考文件[3]，這裡我們使用的本地訓練程式碼如下：

1instance_type='local_gpu'
2instance_count = 1
3CUDA_VISIBLE_DEVICES='0'
4
5hp= {
6    'stage': 4, 'stop_stage': 4, 'train_set':'train', 
7    'data': data_dir, 'trail_dir': trail_dir, 'shared_dir': shared_dir,
8    'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 
9    'num_nodes': instance_count
10}
11
12estimator=PyTorch( 
13    entry_point='examples/aishell/s0/sm-run.sh',
14    image_uri=training_repository_uri,
15    instance_type =instance_type,
16    instance_count=instance_count,
17    source_dir='.',
18    role=role,
19    hyperparameters=hp,
20
21    subnets=subnets,
22    security_group_ids=security_group_ids,
23
24    debugger_hook_config=False,
25    disable_profiler=True
26)
27
28
29estimator.fit({'train': 'file:///fsx'})
30Python

程式碼的輸出如下：

1Creating 2n0im72bz3-algo-1-tpyyu ... 
2Creating 2n0im72bz3-algo-1-tpyyu ... done
3Attaching to 2n0im72bz3-algo-1-tpyyu
4…
52n0im72bz3-algo-1-tpyyu | Invoking script with the following command:
62n0im72bz3-algo-1-tpyyu | 
72n0im72bz3-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
8…
92n0im72bz3-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
102n0im72bz3-algo-1-tpyyu | algo-1-tpyyu: 2021-06-24 15:50:09,669 INFO     [train.py:228] Epoch 0 TRAIN info lr 8e-08
112n0im72bz3-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
122n0im72bz3-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 0Python

上述引數中，source_dir指定的路徑會被打包上傳到Amazon S3，然後，下載到容器例項中。這樣的話，我們每次的程式碼變更都可以直接體現在容器中。

此外，在使用本地訓練模式時，Amazon SageMaker會藉助本地的docker-compose啟動對應的訓練任務，您可以在/tmp目錄下找到相關的docker-compose檔案，

如/tmp/tmp6y009akq，我們可以觀察到如下內容：

1sh-4.2$ tree /tmp/tmp6y009akq
2/tmp/tmp6y009akq
3├── artifacts
4├── docker-compose.yaml
5├── model
6└── output
7    └── data
8Bash

其中，docker-compose.yaml包含了相關的配置資訊，內容如下：

1sh-4.2$ cat /tmp/tmp6y009akq/docker-compose.yaml 
2networks:
3  sagemaker-local:
4    name: sagemaker-local
5services:
6  algo-1-tpyyu:
7    command: train
8    container_name: 2n0im72bz3-algo-1-tpyyu
9    environment:
10    - AWS_REGION=us-east-1
11    - TRAINING_JOB_NAME=sagemaker-wenet-2021-06-24-15-49-58-018
12    image: <your-aws-account-id>.dkr.ecr.us-east-1.amazonaws.com/sagemaker-wenet:training-pip-pt181-py38
13    networks:
14      sagemaker-local:
15        aliases:
16        - algo-1-tpyyu
17    stdin_open: true
18    tty: true
19    volumes:
20    - /tmp/tmp6y009akq/algo-1-tpyyu/output:/opt/ml/output
21    - /tmp/tmp6y009akq/algo-1-tpyyu/output/data:/opt/ml/output/data
22    - /tmp/tmp6y009akq/algo-1-tpyyu/input:/opt/ml/input
23    - /tmp/tmp6y009akq/model:/opt/ml/model
24    - /opt/ml/metadata:/opt/ml/metadata
25    - /fsx:/opt/ml/input/data/train
26version: '2.3'
27Bash

可以看到，docker-compose通過volumes引數，將本地的路徑對映為容器裡的目錄，而不需要執行訓練資料的二次複製。

託管訓練模式

在確定程式碼邏輯無誤後，我們可以很容易通過修改引數的方式，使用託管的例項開啟真正的訓練任務。

這裡，我們只需要調整例項型別、需要的例項數量及資料輸入方式。我們以2臺ml.p3.8xlarge的例項為例，其各自包含4張Tesla V100顯示卡，共8張顯示卡。

訓練程式碼如下：

1instance_type='ml.p3.8xlarge'
2instance_count = 2
3CUDA_VISIBLE_DEVICES='0,1,2,3'
4
5hp= {
6    'stage': 4, 'stop_stage': 4, 'train_set':'train', 
7    'data': data_dir, 'trail_dir': trail_dir, 'shared_dir': shared_dir,
8    'CUDA_VISIBLE_DEVICES': CUDA_VISIBLE_DEVICES, 
9    'ddp_init_protocol': 'tcp',
10    'num_nodes': instance_count
11}
12
13estimator=PyTorch( 
14    entry_point='examples/aishell/s0/sm-run.sh',
15    image_uri=training_repository_uri,
16    instance_type =instance_type,
17    instance_count=instance_count,
18    source_dir='.',
19    role=role,
20    hyperparameters=hp,
21
22    subnets=subnets,
23    security_group_ids=security_group_ids,
24
25    debugger_hook_config=False,
26    disable_profiler=True,
27    environment={
28        'NCCL_SOCKET_IFNAME': 'eth0',
29        'NCCL_IB_DISABLE': 1
30    }
31)
32
33estimator.fit(inputs={'train': file_system_input_train})
34Python

其中，引數CUDA_VISIBLE_DEVICES需設定為訓練例項的GPU卡數量。如果僅有一張GPU顯示卡，則其值為’0’。

這裡需要注意的是，撰寫本文時，Amazon SageMaker訓練任務在掛載FSx時，還不支援指定掛載選項flock，導致無法使用基於file的分散式初始化方法。因此，我們簡單調整WeNet的訓練程式碼，轉而使用基於TCP的初始化方法，來繼續模型訓練。

您還可以觀察到，我們傳入了environment引數，其表示設定容器中對應的環境變數。由於Amazon SageMaker拉起的訓練例項會包含不止一個網路卡，因此，我們需要通過NCCL_SOCKET_IFNAME環境變數，將NCCL使用的網路卡設定為eth0。

此外，Amazon SageMaker支援使用競價例項來訓練模型，以有效降低成本，您可以參考文件[4]檢視使用方法。

模型檔案

在訓練完成之後，會在您設定的目錄生成對應的模型檔案，本文為/fsx/sm-train/trail0/exp/unified_transformer目錄。

如果您需要匯出支援序列化和優化的(TorchScript)模型，則可以調整hp變數中的stage及stop_stage，通過本地模式執行訓練程式碼即可。有關TorchScript，可以參考[5]。

相關程式碼邏輯如下：

1instance_type='local_gpu'
2…
3hp= {
4    'stage': 5, 'stop_stage': 6, 'train_set':'train', 
5…
6}
7
8estimator=PyTorch(
9…
10)
11
12estimator.fit({'train':'file:///fsx'})
13Python

執行完成之後，會在上述目錄生成對應的模型檔案final.zip及量化模型final_quant.zip檔案。

現在，我們已經完成了一次模型訓練工作。我們知道，想要得到一個滿足當下需求的模型，需要經歷多次試驗，多次迭代及訓練。您可以通過上述方法，在Amazon SageMaker上快速嘗試不同的超引數或者其他的演算法，而無需考慮如何配置機器學習的基礎環境等運維相關工作。

模型託管

至此，我們已經得到了訓練好的模型檔案。您可以通過Amazon SageMaker部署模型，也可以通過其他方式部署。在後續的文章中，我們會詳細介紹如何在Amazon Web Services部署訓練好的模型。