对 PyTorch BERT 模型进行微调，并将其部署到 Amazon SageMaker 上的 Amazon Elastic Inference

Original URL: https://aws.amazon.com/cn/blogs/machine-learning/fine-tuning-a-pytorch-bert-model-and-deploying-it-with-amazon-elastic-inference-on-amazon-sagemaker/

文本分类，是一种将不同文本内容划分到对应类别的技术，其拥有广泛的应用范围：电子邮件服务商通过文本分类检测垃圾邮件，营销机构借此对客户评论进行情感分析，论坛版主则借此检测不当发帖等等。

以往，数据科学家使用 tf-idf, word2vec或 bag-of-words (BOW)等方法，生成用于训练分类模型的特征。尽管这些技术在诸多自然语言处理（NLP）任务中获得了巨大成功，但在不同的上下文背景之下，其往往无法准确地捕捉单词含义。最近，随着基于Transformers的双向编码器表示（BERT）技术在结合实际上下文准确实现单词含义编码方面带来的突出表现，人们也希望借助BERT的力量在文本分类任务当中获得更理想的结果。

Amazon SageMaker 是一项全托管服务，能够为开发人员及数据科学家提供快速构建、训练并部署机器学习（ML）模型的能力。Amazon SageMaker消除了ML流程中各个步骤带来的繁重工作，极大降低了高质量模型的开发门槛。Amazon SageMaker Python SDK还提供开源API与容器，允许您更轻松地在Amazon SageMaker中使用多种不同ML与深度学习框架，实现模型的训练与部署作业。

我们的客户经常需要快速调优并轻松部署NLP模型。此外，客户也希望尽可能降低推理延迟与模型推理成本。Amazon Elastic Inference能够将GPU推理加速能力附加至CPU类型的终端端点当中，可以在不牺牲性能的前提下显著降低深度学习的推理成本。

本文将介绍如何使用Amazon SageMaker对PyTorch BERT模型进行微调，并将其部署在应用了Elastic Inference的SageMaker终端节点上。本文中使用的全部代码皆发布在GitHub repo之上。关于BERT微调的更多详细信息，请参阅PyTorch BERT调优教程。

BERT是什么？

BERT最初发布于2018年11月，这是一种革命性的模型，能够主动屏蔽句子中的一个或者多个单词。BERT将屏蔽过单词的句子作为输入，借此自我训练以预测被屏蔽的单词内容。此外，BERT还能够应用于预测下一句的任务。

BERT代表着一项重大突破，已经帮助业界研究人员及数据工程师在众多NLP任务中取得重大成果。BERT提供的各个单词的表征能够切实与所处上下文（即句子中的其余部分）相匹配。关于BERT的更多详细信息，请参阅BERT：用于语言理解的深度双向Transformers预训练模型。

BERT调优

数据科学家在NLP项目当中面临的最大挑战之一，在于缺乏训练数据。大家往往只能获得几千条带有人工标记的文本数据，用于模型训练。但是，现代深度学习NLP任务又需要大量标记数据，而解决此难题的一大重要方法，就是使用迁移学习技术。

迁移学习是一种ML方法，旨在将预训练完成的模型（比如用于图像分类的预训练ResNet模型）重新用作另一不同、但具有相关性的问题。通过复用预训练模型中的参数，我们可以节约大量的训练时间与成本。

BERT是基于BookCorpus与英文维基百科的数据进行训练，二者分别包含8亿以及25亿个单词[1]。从零开始训练BERT的成本极为高昂，但通过迁移学习，大家可以面对新的场景用例时使用相关少量的训练数据对BERT进行快速微调，借此实现常见NLP任务（例如文本分类与问题解答）的高质量预测结果。

解决方案概述

在本文中，我们将分步介绍数据集、训练流程以及最终的模型部署环节。

我们使用Amazon SageMaker notebook实例用于代码运行。关于在Amazon SageMaker上使用Jupyter notebooks的更多详细信息，请参阅使用Amazon SageMaker notebook实例，或者Amazon SageMaker Studio入门指南。

本文中的notebook与代码皆发布于 GitHub之上。您可以克隆 GitHub repo并打开Jupyter notebook文件。

问题与数据集

在本文中，我们使用语言可接受性语料库（CoLA），这是一套对从已出版语言学文献中收集到的10657个英语句子进行符合语法与不符合语法标记的数据集。在我们的notebook中，将使用以下代码下载并解压这些数据：

if not os.path.exists("./cola_public_1.1.zip"):
    !curl -o ./cola_public_1.1.zip https://nyu-mll.github.io/CoLA/cola_public_1.1.zip
if not os.path.exists("./cola_public/"):
    !unzip cola_public_1.1.zip

在训练数据中，我们只需要其中两列——句子本体及其标签：

df = pd.read_csv(
    "./cola_public/raw/in_domain_train.csv",
    sep="\t",
    header=None,
    usecols=[1, 3],
    names=["label", "sentence"],
)
sentences = df.sentence.values
labels = df.label.values

如果我们输出部分句子，即可看到该数据集如何根据句子语法的完整性进行句子标记。具体参见以下代码：

print(sentences[20:25])
print(labels[20:25])

["The professor talked us." "We yelled ourselves hoarse."
 "We yelled ourselves." "We yelled Harry hoarse."
 "Harry coughed himself into a fit."]
[0 1 0 0 1]

接下来，我们对数据集进行拆分以进行训练与测试，而后将其上传至Amazon S3以供后续使用。SageMaker Python SDK可帮助我们快速完成上传操作：

from sagemaker.session import Session
from sklearn.model_selection import train_test_split

train, test = train_test_split(df)
train.to_csv("./cola_public/train.csv", index=False)
test.to_csv("./cola_public/test.csv", index=False)

session = Session()
inputs_train = session.upload_data("./cola_public/train.csv", key_prefix="sagemaker-bert/training/data")
inputs_test = session.upload_data("./cola_public/test.csv", key_prefix="sagemaker-bert/testing/data")

训练脚本

在本文中，我们使用 PyTorch-Transformers库。此库中包含用于BERT等多种NLP模型的PyTorch实现与预训练模型权重。详见以下代码：

model = BertForSequenceClassification.from_pretrained(
    "bert-base-uncased",  # Use the 12-layer BERT model, with an uncased vocab.
    num_labels=2,  # The number of output labels--2 for binary classification.
    output_attentions=False,  # Whether the model returns attentions weights.
    output_hidden_states=False,  # Whether the model returns all hidden-states.
)

根据SageMaker PyTorch镜像的规定，我们的训练脚本应将在训练过程中学习到的模型文件保存至文件路径model_dir。训练完成之后，Amazon SageMaker将保存在model_dir中的模型文件上传至Amazon S3以进行下一步部署。脚本将使用以下代码保存训练得出的模型工件：

model_2_save = model.module if hasattr(model, "module") else model
model_2_save.save_pretrained(save_directory=args.model_dir)

我们将此脚本保存为 train_deploy.py文件，并将该文件放置在名为code/的目录当中。大家可以在该目录中查看完整的训练脚本。

由于PyTorch-Transformer本身并不包含在Amazon SageMaker PyTorch镜像当中，因此我们需要提供对应的 requirements.txt文件，保证Amazon SageMaker能够安装该库以进行训练与推理。 requirements.txt文件属于文本文件，其中包含使用pip install进行安装的条目列表。您也可以指定需要安装的各条目的具体版本。要安装PyTorch-Transformer，我们需要将以下行添加至 requirements.txt文件当中。

transformers==2.3.0

您可以在 GitHub repo上查看完整文件，也可以通过 code/目录进行查看。关于requirements.txt文件的更多详细信息，请参阅Requirements文件。

在Amazon SageMaker上执行训练

我们使用Amazon SageMaker对我们的自定义PyTorch代码执行模型训练与部署。Amazon SageMaker Python SDK能够极大降低在Amazon SageMaker中运行PyTorch脚本的难度。接下来，我们可以使用SageMaker Python SDK对经过训练的模型加以部署，并运行实际预测。关于将SDK与PyTorch配合使用的更多详细信息，请参阅将PyTorch与SageMaker Python SDK配合使用。

首先，我们使用PyTorch estimator进行模型训练。在创建此 estimator 时，请注意明确指定以下内容：

entry_point – PyTorch脚本的名称
source_dir – 训练脚本与 requirements.txt 文件的位置
framework_version: 我们希望使用的PyTorch版本

PyTorch estimator支持多机分布式PyTorch训练。要使用此功能，我们只需将train_instance_count的值设定为大于1即可。我们的训练脚本仅支持面向GPU实例进行分布式训练。

在估计器创建完成之后，我们调用fit()以启动一项训练作业。接下来，我们使用之前上传训练数据时获得的Amazon S3 URI，详见以下代码：

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    entry_point="train_deploy.py",
    source_dir="code",
    role=role,
    framework_version="1.3.1",
    py_version="py3",
    train_instance_count=2,
    train_instance_type="ml.p3.2xlarge",
    hyperparameters={
        "epochs": 1,
        "num_labels": 2,
        "backend": "gloo",
    }
)
estimator.fit({"training": inputs_train, "testing": inputs_test})

在训练开始之后，Amazon SageMaker会显示训练进度（如以下代码所示），具体包括轮次、训练损失以及测试数据精度：

2020-06-10 01:00:41 Starting - Starting the training job...
2020-06-10 01:00:44 Starting - Launching requested ML instances......
2020-06-10 01:02:04 Starting - Preparing the instances for training............
2020-06-10 01:03:48 Downloading - Downloading input data...
2020-06-10 01:04:15 Training - Downloading the training image..
2020-06-10 01:05:03 Training - Training image download completed. Training in progress.
...
Train Epoch: 1 [0/3207 (0%)] Loss: 0.626472
Train Epoch: 1 [350/3207 (98%)] Loss: 0.241283
Average training loss: 0.5248292144022736
Test set: Accuracy: 0.782608695652174
...

我们可以监控训练进度，请保证在继续进行notebook中的后续部分之前，确认训练流程已经成功完成。

部署脚本

在模型训练完成之后，我们通过在PyTorch estimator上调用 deploy 将模型托管在Amazon SageMaker终端节点之上。该终端节点将运行一套Amazon SageMaker PyTorch模型服务器。我们需要对此服务器中的两项组件加以配置：模型加载与模型服务。这两个组件的实现通过推理脚本train_deploy.py完成，完整文件可通过 GitHub repo获取。

model_fn()函数用于加载已保存模型，并返回一个模型对象以供模型服务组件使用。SageMaker PyTorch模型服务器通过调用model_fn加载我们的模型：

def model_fn(model_dir):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = BertForSequenceClassification.from_pretrained(model_dir)
    return model.to(device)

input_fn() 对预测输入进行反序列化与数据转换。在本用例中，我们的请求正文将首先被序列化为JSON格式，而后发送至模型服务端点。接下来，我们首先在 input_fn()中对JSON格式的请求正文进行反序列化，然后根据BERT的要求将输入以 torch.tensor的形式返回：

def input_fn(request_body, request_content_type):
    if request_content_type == "application/json":
        sentence = json.loads(request_body)
    
        input_ids = []
        encoded_sent = tokenizer.encode(sentence,add_special_tokens = True)
        input_ids.append(encoded_sent)
    
        # pad shorter sentences
        input_ids_padded =[]
        for i in input_ids:
            while len(i) < MAX_LEN:
                i.append(0)
            input_ids_padded.append(i)
        input_ids = input_ids_padded
    
        # mask; 0: added, 1: otherwise
        [int(token_id > 0) for token_id in sent] for sent in input_ids

        # convert to PyTorch data types.
        train_inputs = torch.tensor(input_ids)
        train_masks = torch.tensor(attention_masks)
    
        # train_data = TensorDataset(train_inputs, train_masks)
        return train_inputs, train_masks predict_fn() 执行预测并返回结果。详见以下代码：

def predict_fn(input_data, model):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()
    input_id, input_mask = input_data
    input_id.to(device)
    input_mask.to(device)
    with torch.no_grad():
        return model(input_id, token_type_ids=None,attention_mask=input_mask)[0]

预构建的Amazon SageMaker PyTorch镜像中的默认支持对预测结果进行序列化。

部署端点

要部署我们的端点，需要在PyTorch estimator对象上调用deploy()，并提供所需数量的实例与实例类型：

predictor = estimator.deploy(initial_instance_count=1, instance_type="ml.m4.xlarge")

接下来，我们通过配置让预测变量使用"application/json"作为内容类型，而后将请求发送至我们的端点：

from sagemaker.predictor import json_deserializer, json_serializer

predictor.content_type = "application/json"
predictor.accept = "application/json"
predictor.serializer = json_serializer
predictor.deserializer = json_deserializer

最后，我们使用预测变量对象以调用该端点：

result = predictor.predict("Somebody just left - guess who.")
print(np.argmax(result, axis=1))

[1]

预测出的类别为 1，符合我们的预期，因为用于测试的句子确实拥有正确的语法表达。

使用Elastic Inference部署端点

要为推理任务选择正确的实例类型，我们需要在不同数量的GPU、CPU以及内存资源之间做出权衡。在独立GPU实例上针对其中某一种资源进行优化，往往会导致其他资源得不到充分利用。Elastic Inference则能够为特定端点提供适量的GPU驱动型推理加速资源，解决了这一难题。自2020年3月起，用户已经可以在Amazon SageMaker与Amazon EC2上获得Elastic Inference对PyTorch的支持能力。

要使用Elastic Inference，我们需要首先将训练完成的模型转换为TorchScript。关于更多详细信息，请参阅使用Amazon Elastic Inference在Amazon SageMaker for PyTorch模型上降低ML推理成本。

我们首先从Amazon S3处下载训练完成的模型文件。模型文件的位置为estimator.model_data。接下来，我们使用以下代码将模型转换为TorchScript：

model_torchScript = BertForSequenceClassification.from_pretrained("model/", torchscript=True)
device = "cpu"
for_jit_trace_input_ids = [0] * 64
for_jit_trace_attention_masks = [0] * 64
for_jit_trace_input = torch.tensor([for_jit_trace_input_ids])
for_jit_trace_masks = torch.tensor([for_jit_trace_input_ids])

traced_model = torch.jit.trace(
    model_torchScript, [for_jit_trace_input.to(device), for_jit_trace_masks.to(device)]
)
torch.jit.save(traced_model, "traced_bert.pt")

subprocess.call(["tar", "-czvf", "traced_bert.tar.gz", "traced_bert.pt"])

要加载TorchScript模型并将其应用于实际预测，我们还需要对模型的加载与预测函数做出些许调整。我们需要创建一个新的脚本 deploy_ei.py，其内容与train_deploy.py脚本略有不同。

要加载模型，我们使用 torch.jit.load替代之前使用的 BertForSequenceClassification.from_pretrained调用：

loaded_model = torch.jit.load(os.path.join(model_dir, "traced_bert.pt"))

要进行预测，我们在最终return语句当中使用torch.jit.optimized_execution：

with torch.no_grad():
    with torch.jit.optimized_execution(True, {"target_device": "eia:0"}):
        return model(input_id,attention_mask=input_mask)[0]

完整的 deploy_ei.py脚本可通过 GitHub repo获取。使用这套脚本，我们即可通过Elastic Inference进行模型部署：

predictor = pytorch.deploy(
    initial_instance_count=1, 
    instance_type="ml.m5.large",
    accelerator_type="ml.eia2.xlarge"
)

通过使用 accelerator_type="ml.eia2.xlarge"参数，我们即可将Elastic Inference加速器附加至终端节点当中。

资源清理

在实验完成之后，请及时删除期间创建的Amazon SageMaker端点以及Amazon SageMaker notebook实例，以避免产生不必要的费用。具体参见以下代码：

predictor.delete_endpoint()

总结

在本文中，我们使用Amazon SageMaker以BERT为起点，训练出一套能够标记句子语法完整性的模型。接下来，我们将模型分别部署在使用Elastic Inference与不使用Elastic Inference的Amazon SageMaker终端节点。您也可以使用这套解决方案对BERT做其他方向的微调，或者使用PyTorch-Transformers提供的其他预训练模型。关于将PyTorch与Amazon SageMaker配合使用的更多详细信息，请参阅将PyTorch与Amazon SageMaker配合使用。

参考文献

[1] Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba以及Sanja Fidler。2015年。《书籍与电影的映射：在观看电影与阅读书籍中实现相似故事的视觉解释》，IEEE国际计算机视觉会议论文集，第19至27页。

亚马逊AWS官方博客