基于TVM的NLP-BERT模型编译及优化实践

0.引言

TVM是一种从端到端的深度学习编译框架，用于优化深度学习模型在CPU、GPU、ARM等任意目标环境下的推理运行速度，相对于TensorFlow、MXNet、Caffe 和 PyTorch等深度学习框架有以下优点：

• 更适合嵌入式场景；

• 支持客户端-服务器的RPC调用；

• 能通过优化提升模型的Inference性能；

• 容易部署，支持多种语言：JS, Java, Python, C++语言；

• 支持多种模型格式：ONNX，Tensorflow, MXnet模型，DarkNet模型等；

• 支持多种硬件平台及使用环境上的模型Inference：ARM，FPGA，CPU，nVidia GPU，Mali GPU等。

自然语言处理NLP（Natural Language Processing）是人工智能和语言学领域的分支学科，能够挖掘自然语言文本中蕴含的信息和知识。

这里我们选用TVM框架去进行编译，此次编译的模型为NLP中的BERT模型，任务类型为问答任务。目标是在文本中找出答案的起始位置 (Start，End)。如下图所示，将 Question 和 Paragraph 传入 BERT，然后 BERT 根据 Paragraph 所有单词的输出预测 Start 和 End 的位置。本文主要是验证TVM框架编译BERT模型进行推理的可行性并进一步对其性能进行优化处理。

1. 数据集预处理

这里使用的是SQuAD（The Stanford Question Answering Dataset 1.1 ）数据集[3]，数据格式如下所示，它整体是一个json格式的数据，"data"中存放的是一个列表，其中存放多个"title"和"paragraphs"组合代表着每一篇文章的元素,每一个"paragraphs"又对应着一个列表，其中包含段落中的一小段描述("context")和若干个问答对("qas")，我们需要从数据集中提取出对应的文本信息、问题以及对应的答案和id信息。预处理部分的代码主要来自于参考文献[2]，下面介绍一下数据处理过程以及主要函数的作用，由于代码篇幅较长，本文不贴具体代码，有兴趣的读者可以自行查看。

{
    "data": [
        {
            "title": "Super_Bowl_50",
            "paragraphs": [
                {
                    "context": "Super Bowl 50 was an American football game to.....",
                    "qas": [
                        {
                            "answers": [
                                {
                                    "answer_start": 177,
                                    "text": "Denver Broncos"
                                }
                            ],
                            "question": "Which NFL team represented the AFC at Super Bowl 50?",
                            "id": "56be4db0acb8001400a502ec"
                        },
                    ...
                },
        },
        {
            "title": "Warsaw",
            "paragraphs": [...
            ]
        },
        {
            "title": "Normans",
            "paragraphs": [...
            ]

        },
        ...

首先遍历整个json数据集进行数据读取，在char_to_word_offset中记录每个字符在"context"字段中的索引位置，这样就可以根据"answer_start"字段的字符从char_to_word_offset中得到问题答案在上下文中的起止位置。最后该函数会返回一个如下的二维列表：

doc_tokens:['Super', 'Bowl', '50', 'was', 'an', 'American', 'football', 'game', 'to', 'determine', 'the', 'champion', 'of', 'the', ...]
end_position:29
orig_answer_text:'Denver Broncos'
qas_id:'56be4db0acb8001400a502ec'
question_text:'Which NFL team represented the AFC at Super Bowl 50?'
start_position:28

_improve_answer_span函数对上边的得到的答案的起止位置进行了修正。

然后convert_examples_to_features函数将数据转换为特征值，在处理过程中问题若超过max_query_length则会截断取前半部分，文档若超过max_seq_length则会使用滑窗法，_check_is_max_context函数会计算每个token在每一段滑窗中的最佳位置，其中重点关注的三个特征值如下:

• input_ids : tokenizer将tokens映射成的ids数组。

• input_mask : 由0，1组成的数组，为了标识是否为补位，如果只有一句话，就用1表示，0用来补位。

• segment_ids : 由0，1组成的数组，为了标识是否为同一个segment（比如query的segment_ids为0，answer的segment_ids为1)。

举例如下：

input_ids : [101, 2040, 2003, 1996, 2132, 2873, 1997, 1996, 14169, 1029, 102, 2206, 2037, 3279, ...]
input_mask : [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...]
segment_ids : [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, ...]

2. BERT模型

本文使用的模型是pytorch jit类型，固定了BERT模型的banchsize为64，模型的input_shape为[64,384]，该模型的输入输出结构如下图：

图中可以看出，BERT模型与一般模型不同，它有三个输入以及两个输出。输入的 input_ids，attention_mask，token_type_ids 分别对应第一节中讲到的特征值 input_ids，input_mask，segment_ids。而输出的 2066 以及 2068 分别对应的是问答任务中答案在上下文中的起始位置和终止位置。

3. TVM推理过程

（1）模型导入

首先利用torch.jit.load 加载 bert.jit 模型，接下来使用 TVM 的 Relay 将 pytorch 模型变成 TVM 可以识别的 Graph IR，TVM 在 Relay 中提供了一个frontend.from_pytorch 用来加载 pytorch 模型并转换为 Relay IR。其中 self.mod 中包含模型中所有输入 Tensor 的 shape 信息, 比如卷积层的 weight 和 bias。self.params 则保存了模型中所有 OP 的权重信息，以字典形式存放， key 是权重 Tensor 的名字。relay.build 构造了 Relay 的计算图, 生成了 self.lib, 然后使用 graph_executor.GraphModule 将 lib 包装在 GraphExecutor 中，创建了一个可以直接从 Python 调用的模块。具体代码如下:

def load_model(self):
    """Load torch model and build module,self.input_shape=[64,384]"""
    self.model_path="/model/bert.jit"
    self.jit_model = torch.jit.load(self.model_path, map_location="cpu")
    shape_list = [
        ("input_ids", self.input_shape),
        ("input_mask", self.input_shape),
        ("segment_ids", self.input_shape),
    ]
    self.mod, self.params = relay.frontend.from_pytorch(self.jit_model, shape_list)
    if self.transform:
        self._pre_process()

    with tvm.transform.PassContext(opt_level=3):
        self.lib = relay.build(self.mod, target=self.target, params=self.params)
    dev = tvm.device(str(self.target), 0)
    self.module = graph_executor.GraphModule(self.lib["default"](dev))

（2）推理结果

从第一节中预处理后的数据中取出模型需要的三个输入输送到self.model，利用set_input接口进行推理，get_output输出结果，结果为两个shape均为[64,384]的NDArray，转成numpy格式取最大值找出答案所对应的可能性最大的初始位置和结束位置，由于BERT模型多输入且size相同的特殊性，在进行推理的时候需要注意输入数据的顺序以及名称，如果TVM中导入pytorch模型时指定的数据名称和最终推理时set_input的数据名称不一致会导致TVM无法得到正确的数据而出现推理出随机结果的情况，如果顺序不对会导致推理错误但不直接报错的情况，不易debug，需要注意，具体代码如下:

data_set = PreSquad()
for i in range(self.input_shape[0]):
    self.input_ids_list.append(data_set.eval_features[i].input_ids)
    self.segment_ids_list.append(data_set.eval_features[i].segment_ids)
    self.input_mask_list.append(data_set.eval_features[i].input_mask)
self.module.set_input(
            input_ids=tvm.nd.array(np.array(self.input_ids_list).astype("float32")),
            input_mask=tvm.nd.array(np.array(self.input_mask_list).astype("float32")),
            segment_ids=tvm.nd.array(np.array(self.segment_ids_list).astype("float32")),
        )
self.module.run()
start_tvm = np.argmax(self.module.get_output(0).numpy(), axis=-1)
end_tvm = np.argmax(self.module.get_output(1).numpy(), axis=-1)

为了对比TVM推理是否准确，本文还利用torch进行了推理,代码如下:

with torch.no_grad():
    torch_start_out, torch_end_out = self.jit_model(
        torch.tensor(self.input_ids_list),
        torch.tensor(self.input_mask_list),
        torch.tensor(self.segment_ids_list),
    )
    start_torch = np.argmax(torch_start_out.numpy(), axis=-1)
    end_torch = np.argmax(torch_end_out.numpy(), axis=-1)

TVM和torch推理结果一致，结果如下：

TVM start_output:
 [ 64 253 133 254 346  57  75 134 170 210  70  98 143 174 299  46 113  66
 220 349  12 281 108  22  28 318 282   0  17  42 240  30  32  41  85  98
 195  21   0 101   0 304 176 117   0  25   0  12   0  25   0 216  88 304
 176   0 283  93   0 204  76 204 346  44]
TVM end_output:
 [ 67 255 134 256 347  57  78 135 170 224  73 252 148 174 301  47 114  69
 220 350  15 283 110  23  30 328 284   0  17  39 243  47  34  41  86  98
 202  30   0 109   0 306 178 118   0  34   0  13   0  34   0 216  88 306
 178   0 287  97   0 211   0 219 348  52]
Pytorch start_output:
 [ 64 253 133 254 346  57  75 134 170 210  70  98 143 174 299  46 113  66
 220 349  12 281 108  22  28 318 282   0  17  42 240  30  32  41  85  98
 195  21   0 101   0 304 176 117   0  25   0  12   0  25   0 216  88 304
 176   0 283  93   0 204  76 204 346  44]
Pytorch end_output:
 [ 67 255 134 256 347  57  78 135 170 224  73 252 148 174 301  47 114  69
 220 350  15 283 110  23  30 328 284   0  17  39 243  47  34  41  86  98
 202  30   0 109   0 306 178 118   0  34   0  13   0  34   0 216  88 306
 178   0 287  97   0 211   0 219 348  52]

4. 性能优化

在上面的实验中我们已经验证了TVM编译BERT模型进行推理的正确性，接下来本文就时延性能利用TVM的AutoScheduler模块对BERT模型推理速度进行优化。AutoScheduler（别名Ansor）指的是无模板的自动调优模块。它将网络划分为小的子图，并对它们进行独立调整。任务调度器对时间进行切分，并将时间资源进行动态分配，最终获取最优的调度方案。

为了优化整个BERT网络模型，我们首先需要调用extract_tasks()函数从之前得到self.mod和self.params对任务进行抽取。extract_tasks()函数会返回任务的key和相应的权重并创建SearchTask对象。

"""Tune kernel of graph in autoschedule"""
print("Extract tasks...")
tasks, task_weights = auto_scheduler.extract_tasks(
    self.mod["main"], target=self.target, params=self.params
)
for idx, task in enumerate(tasks):
    print(f"========== Task {idx} (workload key: {task.workload_key}) ==========")
    print(task.compute_dag)

对于BERT模型，其输出如下:

之后进行tuning，函数auto_scheduler.TaskScheduler()可以设置一些相关的参数，num_measure_trials代表的是测试的次数，搜过总共会进行num_measure_trials个调度并返回其中最好的一个。在试运行阶段可以将其设置为较小的数量，一般情况下会按照tasks*800去进行设置。该参数很大程度会影响整个搜索的时间。runner运行程序并测量时间成本。measure_callbacks是搜索回调函数，会将测试记录转存到我们指定的文件中，方便保存和查询。

search_parameter = 800
tuner = auto_scheduler.TaskScheduler(tasks, task_weights)
tune_option = auto_scheduler.TuningOptions(
    num_measure_trials=search_parameter*len(tasks),
    runner=auto_scheduler.LocalRunner(repeat=10, enable_cpu_cache_flush=True, timeout=50),
    measure_callbacks=[auto_scheduler.RecordToFile(self.auto_scheduler_file)],
)
tuner.tune(tune_option)

下表是搜索中生成的表格，其中可以观察到所使用的的BERT模型具体的task名称、任务时延和速度、测试次数和下一个任务ID等。另外搜索过程中也会出现一些无效调度导致的错误，这不会影响到整个过程，无需关注。

搜索完成后，self.auto_scheduler_file文件中会储存最优的调度计划，利用该文件重新进行编译和评估，代码如下:

print("Compile...")
with auto_scheduler.ApplyHistoryBest(self.auto_scheduler_file):
    with tvm.transform.PassContext(
        opt_level=3, config={"relay.backend.use_auto_scheduler": True}
    ):
        self.lib = relay.build(self.mod, target=self.target, params=self.params)
# Create graph executor
dev = tvm.device(str(self.target), 0)
module = graph_executor.GraphModule(self.lib["default"](dev))
self.module.set_input(
    input_ids=tvm.nd.array((np.random.uniform(size=self.input_shape)).astype("float32")),
    input_mask=tvm.nd.array((np.random.uniform(size=self.input_shape)).astype("float32")),
    segment_ids=tvm.nd.array((np.random.uniform(size=self.input_shape)).astype("float32")),
)
# Evaluate
print("Evaluate inference time cost...")
print(module.benchmark(dev, repeat=3, min_repeat_ms=500))

TVM优化前后的BERT模型推理时间如下，可以看出利用Ansor优化后，BERT模型的时延性能明显提升。

Evaluation of the PyTorch network without auto tune:
 mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  
 27109.7505   27113.5095   27168.2594   27047.4827    49.3784 
Evaluation of the PyTorch network with auto tune:
 mean (ms)   median (ms)    max (ms)     min (ms)     std (ms)  
 12501.2272   12484.2228   12580.4266   12439.0323    58.9630

5. 总结

本文利用TVM编译框架对NLP中的BERT模型针对问答任务进行了推理并简述了squad数据集的预处理过程，另外利用torch框架辅助验证了TVM推理的正确性。在推理性能方面，描述了TVM的AutoScheduler模块对BERT网络模型优化的过程以及效果。

参考文献

1.BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding Attention Is All You Need  

2.https://github.com/mlcommons/inference/blob/f5367250115ad4febf1334b34881ab74f2e55bfe/language/bert  

3.https://rajpurkar.github.io/SQuAD-explorer/  

4.https://tvm.apache.org/docs/how_to/tune_with_autoscheduler/  tune_network_x86html#sphx-glr-how-to-tune-with-autoscheduler-tune-network-x86-py  

5.https://tvm.apache.org/docs/reference/api/python/auto_scheduler.html