使用ONNX和Torchscript加快推理速度的测试

2021-03-10 14:40:16 浏览数 (1)

近年来,基于Transformer 架构的模型一直是推动NLP在研究和工业上取得突破的动力。BERT、XLNET、GPT或XLM是一些改进了技术水平的模型,它们达到了GLUE等流行基准的顶级水平。

这些进步带来了高昂的计算成本,大多数基于Transformer的模型都是庞大的,用于训练的参数数量和数据都在不断增加。最初的BERT模型已经有1.1亿个参数,而最后的GPT-3有1750亿个参数,这是在两年的研究中惊人的增长了1700倍。

这些庞大的模型通常需要数百个GPU进行数天的训练才能发挥作用,幸运的是,多亏了迁移学习,我们可以下载预训练的模型,并在我们自己的更小的数据集上快速地以低成本调整它们。

也就是说,一旦训练完成,我们手中有一个庞大的模型,如果想要将其部署到生产中与其他模型相比,推理需要相对较长的时间,而且它可能太慢,无法达到需要的吞吐量。

虽然您可以投资更快的硬件或使用更多的服务器来完成工作,但有还有其他的方法来减少模型的推理时间:

模型修剪:减少层数,嵌入的尺寸或隐藏层中的单位数。

量化:不使用32位浮点数(FP32)进行加权,而是使用半精度(FP16)甚至8位整数。

将模型从本机Pytorch / Tensorflow导出到适当的格式或推理引擎(Torchscript / ONNX / TensorRT ...)

第一种和第二种方法通常意味着对模型进行重新训练,而后两种方法则是在训练后完成的,本质上与您的特定任务无关。

如果推理速度对用例极为重要,那么很可能需要尝试所有这些方法以生成可靠且快速的模型。但是,在大多数情况下,将模型导出为适当的格式/框架并进行批量预测将为您提供更快的结果,而所需的工作量却最少。在这里,我们将重点介绍这种方法,以了解其对模型吞吐量的影响。

我们将通过一些实验探讨更改模型格式和批处理的影响:

  1. 使用常规的Pytorch CPU / GPU的基准
  2. 将Pytorch模型导出到Torchscript CPU / GPU
  3. 将Pytorch模型转换为ONNX CPU / GPU

所有实验均以1/2/4/8/16/32/64批次运行

截至本文发布时,由于缺乏Pytorch嵌入所使用的int64的支持,因此尚无法将Transformer模型直接从Pytorch导出到TensorRT,因此我们暂时将其跳过。

我们将对Roberta的法语变体camemBERT(〜100M参数)执行句子分类。由于绝大多数计算是在Transformer模型内部完成的,因此无论您执行什么任务,都应得到相似的结果。

首先,我们将快速了解如何将Pytorch模型导出为相关的格式/框架,如果您不想阅读代码,则可以跳至结果部分。

如何导出模型

常规pytorch

尽管有不同的处理方法,但在Pytorch中保存和加载模型非常简单。为了进行推理,官方文档建议保存模型的“ state_dict”,这是一个包含模型可学习参数的python字典。这比保存整个模型进行更轻巧。

代码语言:javascript复制
 #saving
 model = SequenceClassifier()
 train_model(model)
 torch.save(model.state_dict(), 'pytorch_model.pt')
 
 #loading
 model = SequenceClassifier()
 model.load_state_dict(torch.load(PATH))
 model.eval() #Set dropout and batch normalization layers to evaluation mode
 with torch.go_grad():
    logits = model(**batch_x)

Torchscript JIT

TorchScript是一种从Pytorch代码创建可序列化和可优化模型的方法。导出到Torchscript后,你的模型就可以在Python和c 中运行了。

Trace:输入通过模型发送,所有操作都记录在一个将定义您的torchscript模型的图中。

Script:如果您的模型更复杂并且具有诸如条件语句之类的控制流,脚本将检查模型的源代码并将其编译为TorchScript代码。

请注意,由于模型将被序列化,所以在保存模型之后它将不能修改,因此应该将它置于评估(eval)模式中,并在保存之前将它导出到适当的设备上。

如果要在CPU和GPU上进行推理,则需要保存2种不同的模型。

代码语言:javascript复制
 #saving
 jit_sample = (batch_x['input_ids'].int().to(device), batch_x['attention_mask'].int().to(device))
 
 model.eval()
 model.to(device)
 module = torch.jit.trace(model, jit_sample)
 torch.jit.save('model_jit.pt')
 
 #loading
 model = torch.jit.load('model_jit.pt', map_location=torch.device(device))
 logits = model(**batch_x)

ONNX

ONNX为AI模型提供了一种开源格式,大多数框架都可以将它们的模型导出为ONNX格式。除了框架之间的互操作性之外,ONNX还提供了一些优化,可以加速推理。

导出到ONNX稍微复杂一些,但是Pytorch确实提供了一个直接的导出函数,你只需要提供一些关键信息。

opset_version,每个版本都支持一组运算符,一些具有奇特架构的模型可能还不能导出。

input_names和output_names是分配给图的输入和输出节点的名称。

dynamic_axes参数是一个字典,它指示输入和输出变量的哪个维度可能会改变,例如batch_size或序列的长度。

代码语言:javascript复制
 #saving
 input_x = jit_sample ## taking sample from previous example
 
 torch.onnx.export(model, input_x,'model_onnx.pt',export_params=True,  opset_version=11, do_constant_folding=True, input_names = ['input_ids', 'attention_mask'], output_names = ['output'],
 dynamic_axes= {
 'input_ids' : {0 : 'batch_size', 1:'length'},'attention_mask' : {0 : 'batch_size', 1:'length'},
 'output' : {0 : 'batch_size'}
 })
 
 #loading
 model = onnxruntime.InferenceSession(model_onnx)
 batch_x = {
 'input_ids':sample['input_ids'].cpu().numpy(),
 "attention_mask":sample['attention_mask'].cpu().numpy()
 }
 logits = model.run(None, batch_x)

ONNX运行时可以与GPU一起使用,尽管它需要特定版本的CUDA, cuDNN和OS,这使得安装过程在一开始很有挑战性。

实验结果

每种配置都在一个包含1k个不同长度句子的数据集上运行了5次。我们用torch 1.7.1和ONNX 1.6.0测试了2种不同的流行GPU: T4和V100。请记住,结果会随着特定的硬件、包版本和数据集而变化。

推理时间的范围从平均每个样本约50 ms到数据集上的0.6 ms,这取决于硬件设置。

在CPU上,对于batch_size <32, ONNX格式是一个明显的赢家,此时格式似乎不再重要了。如果我们一个样本一个样本地预测,我们会发现ONNX能够以一小部分的成本在GPU上进行基线推理。

正如预期的那样,推理在GPU上要快得多,特别是在批处理大小较大的情况下。我们还可以看到,理想的批处理大小取决于使用的GPU:

  • 对于T4来说,最好的设置是用8个批次的样本运行ONNX,这比pytorch上的批大小为1的速度快了大约12倍
  • 对于批量为32或64的V100,与GPU的基线相比,我们可以达到 28倍的加速,与CPU的基线相比,可以达到 90倍的加速。

总的来说,我们发现选择合适的格式对于较小的批数有显著的影响,但是随着批数的增加,这种影响会缩小,在64批样品中,3种设置之间的差异在10%以内。

序列长度和批处理策略的影响

另一件需要考虑的事情是序列长度。Transformer通常被限制为512个标记的序列,但在这个范围内,不同序列长度的速度和内存需求存在巨大差异。

对于更大的批次,推理时间大致随序列长度线性增加,但对于单个样品不是这样。这意味着,如果您的数据是由长序列的文本(例如新闻文章)组成的,那么通过批处理就不会得到那么大的加速。和往常一样,这取决于您的硬件,V100比T4快,并且在预测长序列时不会受到太多的影响,然而另一方面,我们的CPU确实会完全不知所措:

如果你的数据在纵向上是不同的,而你处理的是批次,这些差异将会导致问题,因为你需要将你的样品填充到批次中最长的样品中,这增加了大量的计算量。因此,通常最好将长度相似的样品批在一起,因为预测多个长度相似的批次比预测一个主要是填充令牌的大批次更容易。

作为测试,让我们看看在运行推理之前对数据集排序时会发生什么:

正如我们所预期的那样,为了获得更大的批量,将相似长度的样品分组在一起有很大的好处。对于未排序的数据,随着批量变大,最终得到一些更长的样本的可能性越来越大,这将显著增加整个批量的推理时间。我们可以看到,从16到64 batch_size降低了20%的推理速度,而排序数据的推理速度提高了10%。

这个策略也可以用来显著地减少训练时间,但是这应该谨慎地做,因为它可能会对您的模型的性能产生负面影响,特别是当标签和样本长度之间存在一些相关性时。

下一个步

虽然这些实验已经直接在Python中运行,但Torchscript和ONNX模型都可以直接在c 中加载,这可以提供额外的推理速度提升。

如果模型对于用例来说仍然太慢,Pytorch提供了不同的量化选项。“动态量化”可以在训练后进行,但它很可能会对模型的准确性产生影响,而“quantization aware training”需要再训练,但它对模型性能的影响应该较小。

总结

正如我们所看到的,没有直接的答案来优化推理时间,因为它主要取决于特定硬件和试图解决的问题。因此应该使用自己的目标硬件和数据进行实验,以获得可靠的结果。

尽管如此,还是有一些很容易实施的指导方针:

预测批次可以提供显著的加速到一定的尺寸(取决于您的特定硬件),特别是如果可将批相似长度的样品放在一起时。

使用Torchscript或ONNX确实为较小的批大小和序列长度提供了显著的加速,在对单个样本运行推理时效果特别强。

ONNX似乎是我们测试过的三种配置中表现最好的,尽管它也是最难安装到GPU上的推理。

Torchscript确实为小批量提供了可靠的加速,而且非常容易设置。

作者:Maxence Alluin

原文地址:https://towardsdatascience.com/an-empirical-approach-to-speedup-your-bert-inference-with-onnx-torchscript-91da336b3a41

deephub翻译组

0 人点赞