使用ONNX将GPT Neo（或其他）投入生产

完整的笔记本：

https://github.com/oborchers/Medium_Repo/blob/master/Putting%20GPT-Neo%20into%20Production%20using%C2%A0ONNX/ONNX-Export.ipynb

介绍

使用Transformer已成为最先进的NLP应用程序的新规范。考虑到BERT或GPT3，我们可以安全地得出结论，几乎所有NLP应用程序都从类似于Transformer的模型中获益匪浅。

然而，这些模型通常部署成本很高，并且需要特殊的硬件来运行。在本文中，你将了解什么是ONNX，以及如何将torch和tensorflow transformers模型移植到ONNX。

你还将学习如何定制torch实现以及如何在之后导出它。具体来说，我们将研究：

1. bert-base-nli-stsb-mean-tokens的简单导出
2. bert-base-nli-stsb-mean-tokens的定制导出
3. 用ORT CustomOps导出Universal Sentence Encoder
4. 尝试导出带有1.3B参数的GPT Neo

什么是ONNX

当我不久前开始使用Transformer的时候，我第一次体验了BERT-as-a-Service。虽然BaaS仍然是一个不错的库，但现在在GPU上部署自己的模型并为其提供一个小型restapi相当简单。通常，这将由一个或多个框架完成，例如torch或tensorflow。但这在实践中有着严重的局限性。

这就是ONNX发挥作用的地方。开放式神经网络交换的目标是提供不同参与者之间的互操作性。互操作性是指：

1. 跨框架共享模型（例如，torch到tensorflow）
2. 跨各种硬件（如CPU、GPU、FPGA等）共享模型

这对社区有好处。尝试在同一GPU上使用两个不同的框架部署模型。这是一种痛苦。

在后台，ONNX提供了一种定制的文件格式，一种由节点组成的计算图，节点本身由基本操作符组成。ONNX拥有大量与深度学习和机器学习相关的核心操作，还提供了使用定制操作的能力。引用他们的主页：https://onnx.ai/about.html

ONNX提供了可扩展计算图模型的定义，以及内置操作符和标准数据类型的定义。 每个计算数据流图都被构造为一个节点列表，这些节点构成一个非循环图。节点有一个或多个输入和一个或多个输出。每个节点会调用某些操作。这个图还有元数据来帮助记录它的目的、作者等。

如果你想了解更多关于ONNX的信息，这里有一个来自微软和NVIDIA的非常好的演示，你可以在这里找到：https://developer.nvidia.com/gtc/2019/video/s9979。请记住，本演示文稿是从2019年开始的，2年内可能会有很多变化。

在开始使用ONNX之前，有三个与我们的目的相关的主要组件：

• ONNX：提供图形格式和操作定义
• ONNX Runtime：提供可用于在硬件上部署模型以进行推断的运行时环境。它包含ExecutionProviders，这使我们能够使用各种方法（如CPU、Cuda或TensorRT）加速操作。
• ONNX Runtime Tools：提供对已转换的ONNX transformers模型执行额外优化的功能。我们不会在这里使用它，但请记住它是存在的！

预备工作

接下来，你将需要许多库。我建议你在继续之前建立自己的Docker映像，它支持最新的NVIDIA驱动程序，甚至可能支持TensorRT。

从NVIDIAs的http://nvcr.io/nvidia/tensorrt:20.12-py3是个好主意。你甚至可能希望从头开始构建ONNXRuntime（推荐）。这是我的Dockerfile文件，https://github.com/oborchers/Medium_Repo/blob/master/onnxruntime-issues/Dockerfile。

此脚本需要根据你的配置进行调整，可能不适合你。它已经在装有V100的容器上进行了测试。这个构建允许你从ONNX运行时访问CPU、CUDA、TensorRT执行提供程序。我们还使用了Transformer库的最新开发版本，即4.5.0.dev0来访问GPT Neo。

1.简单导出

注意：这里提供完整的笔记本:https://github.com/oborchers/Medium_Repo/blob/master/Putting%20GPT-Neo%20into%20Production%20using%C2%A0ONNX/ONNX-Export.ipynb

我们要看的第一个模型是来自句子Transformer库的bert-base-nli-stsb-mean-tokensmodel。该模型也可在hub上使用。它本质上是一个BERT模型，经过训练可以产生良好的句子嵌入，在相似性搜索中也有很好的表现。

为了转换模型，让我们使用transformers库中的convert_graph_to_onnx方法（参见这里）。导出代码如下所示：

# 导出transformers模型的脚本
model_name = "sentence-transformers/bert-base-nli-stsb-mean-tokens"
pipeline_name = "feature-extraction"
model_pth = Path(f"encoder/{model_name}.onnx")

nlp = transformers.pipeline(pipeline_name, model=model_name, tokenizer=model_name, device=0)
tokenizer = nlp.tokenizer

if model_pth.exists():
    model_pth.unlink()

convert_graph_to_onnx.convert(
    framework="pt",
    model=model_name,
    output=model_pth,
    opset=12,
    tokenizer=model_name,
    use_external_format= False,
    pipeline_name= pipeline_name,
)

接下来，我们只需要加载模型，创建一个推理会话。此外，我们传递一些会话选项，并加载导出的模型：

# 我们开始只与CUDA合作
ONNX_PROVIDERS = ["CUDAExecutionProvider", "CPUExecutionProvider"]
opt = rt.SessionOptions()
sess = rt.InferenceSession(str(model_pth), opt, providers=ONNX_PROVIDERS)

model_input = tokenizer.encode_plus(span)
model_input = {name : np.atleast_2d(value) for name, value in model_input.items()}
onnx_result = sess.run(None, model_input)

print(onnx_result[0].shape)
print(onnx_result[1].shape)

看起来不错！模型正在加载，一切都很好。

如果我们比较一下速度，来自transformers的nlp管道在span="Hello my friends!"大约运行10毫秒。这模拟了在线推理，这可能是最常见的用例。另一方面，ONNX模型的运行速度是2.8ms，快了2.5倍，而且只需要几行代码，没有进一步的优化。

理论上，你现在可以从ONNX运行时工具将模型放到前面提到的优化器中。但是要注意：如果你使用use_gpu=True运行优化器，那么请确保安装了不带TensorRT的ONNX运行时，因为如果启用了TensorRT执行提供程序，则导出将不起作用。

如果你仔细看，你可以看到打印声明中产生的形状是不正确的。返回的是两个数组的列表，它们的形状分别是(1,6,768)和(1,768)。理论上，我们期望返回的形状是(1,768)，因为我们使用的是一个句子编码器。

这种行为是由于句子转换器库需要一个额外的平均池层添加到token嵌入之上的管道中。也就是说，如果我们想要一个统一的部署框架，并且不想在之后摆弄numpy或torch，那么在之后添加层并不是一个优雅的解决方案，这会破坏我们的目的。

在我们检查自定义输出之前，让我们先看看基准:

• SENTENCECUDATransformer：12.3 ms± 1.4 ms
• ONNX CUDA（V100）：2.21 ms ± 77 µs
• ONNX TensorRT（V100，ExecutionProvider）：3.86 ms ± 181 µ

坦白说，我们在这里看到的结果很奇怪。我已经在这里打开了一个问题，因为我无法从TensorRT获得任何加速,https://github.com/microsoft/onnxruntime/issues/7230。

2.自定义导出

添加自定义层需要我们了解所使用的convert函数内部的情况。

Spoiler：它相当简单。convert函数调用两个函数，即infer_shapes和ensure_valid_input。然后，所有推断出的形状加上生成的torch.nn.Module对象被传递给torch.onnx.export函数。该文档提供了一个关于如何正确使用导出函数的非常好的示例。，https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html。

理解导出函数最重要的是以下参数：

1. 输入名称：底层torch模型的forward函数的参数。必须按正确的顺序。
2. 输出层名称：输出层的名称。
3. 动态轴：定义哪些轴是动态的，以何种方式是动态的（在未来会更有意义）。
4. 参数：一组通过模型的示例输入。

让我们把它们封装起来：

def print_transformers_shape_inference(name_or_path: str):
    """打印onnx的transformers形状推断。""" 
    res = {}

    model_pipeline = transformers.FeatureExtractionPipeline(
        model=transformers.AutoModel.from_pretrained(name_or_path),
        tokenizer=transformers.AutoTokenizer.from_pretrained(
            name_or_path, use_fast=True
        ),
        framework="pt",
        device=-1,
    )

    with torch.no_grad():
        (
            input_names,
            output_names,
            dynamic_axes,
            tokens,
        ) = convert_graph_to_onnx.infer_shapes(model_pipeline, "pt")
        ordered_input_names, model_args = convert_graph_to_onnx.ensure_valid_input(
            model_pipeline.model, tokens, input_names
        )

    res["input_names"] = input_names
    res["output_names"] = output_names
    res["dynamic_axes"] = dynamic_axes
    res["tokens"] = tokens
    res["exemplary_input"] = model_args

    print()
    print(f"Inferred shapes for {name_or_path}")
    print(f"Input names: {input_names}")
    print(f"Output names: {output_names}")
    print(f"Dynamic Axes:\n{json.dumps(dynamic_axes,sort_keys=True, indent=4)}")
    print(f"Tokens:{tokens}")
    print(f"Ordered input names: {ordered_input_names}")
    print(f"Arguments: {model_args}")

    return res

model_args = print_transformers_shape_inference(model_name)

将打印形状推理应用于感兴趣的BERT模型，我们得到以下形状：

Inferred shapes for sentence-transformers/bert-base-nli-stsb-mean-tokens
Input names: ['input_ids', 'token_type_ids', 'attention_mask']
Output names: ['output_0', 'output_1']
Dynamic Axes:
{
    "attention_mask": {
        "0": "batch",
        "1": "sequence"
    },
    "input_ids": {
        "0": "batch",
        "1": "sequence"
    },
    "output_0": {
        "0": "batch",
        "1": "sequence"
    },
    "output_1": {
        "0": "batch"
    },
    "token_type_ids": {
        "0": "batch",
        "1": "sequence"
    }
}
Tokens:{
    'input_ids': tensor([[ 101, 2023, 2003, 1037, 7099, 6434,  102]]), 
    'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0]]),  
    'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1]])
}
Ordered input names: ['input_ids', 'attention_mask', 'token_type_ids']
Arguments: (
    tensor([[ 101, 2023, 2003, 1037, 7099, 6434,  102]]), 
    tensor([[1, 1, 1, 1, 1, 1, 1]]), 
    tensor([[0, 0, 0, 0, 0, 0, 0]])
)

这是完全有道理的解释。output_0代表pooler_output，output_1代表返回的BaseModelOutputWithPoolingAndCrossAttention的last_hidden_state。input_ids、token_type_ids和attention_mask都是动态的，是tokenizer函数的输出。

让我们继续建立一个简单的torch模型，它继承了BERT模型。我们添加的唯一内容是对token嵌入进行加权求和。

class SentenceTransformer(transformers.BertModel):
    def __init__(self, config):
        super().__init__(config)
        # 为ONNX输出规范命名别名使其更容易识别层
        self.sentence_embedding = torch.nn.Identity()

    def forward(self, input_ids, token_type_ids, attention_mask):
        # 从基本模型中获取token嵌入
        token_embeddings = super().forward(
            input_ids, 
            attention_mask=attention_mask, 
            token_type_ids=token_type_ids
        )[0]
        # 将池化层叠加在其之上
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size())
        sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
        sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
        return self.sentence_embedding(sum_embeddings / sum_mask)

# 基于原始管道的配置创建新的模型
model = SentenceTransformer(config=nlp.model.config).from_pretrained(model_name)

最后检查模型产生的输出与原始模型大致相同，我们可以继续。

在导出我们的新模型之前，唯一要做的就是修改我们之前导出的动态轴和输出名称。这是因为我们现在有了一个不同的输出层，它也是动态的（在批大小上）。我们可以使用标识层的名称来更好地标识输出层。

del model_args["dynamic_axes"]["output_0"] # 删除未使用的输出
del model_args["dynamic_axes"]["output_1"] # 删除未使用的输出
model_args["dynamic_axes"]["sentence_embedding"] = {0: "batch"}

model_args["output_names"] = ["sentence_embedding"]

太好了！现在我们已经准备好了新的ONNX模型，并且可以用它进行推理。输出形状现在是预期的（1768），它几乎等于原始模型。此外，新的模型运行在2.4ms，所以我们没有失去任何速度，并获得了一个适当的端到端模型。

很明显，这个过程可以根据你的喜好定制。还可以在此基础上训练自己的分类器，并以相同的方式将其添加到编码器中。

我们已经创建了前两个ONNX模型。干得好！让我们做点不同的事。

3.使用ORT CustomOps导出

这一部分特别关注universal sentence encoder 5，这是一个我一直在使用的模型，我非常喜欢。速度快，性能好，体积小。

谢天谢地，存在tf2onnx库。tf2onnx是一个导出工具，用于从tensorflow模型生成ONNX文件。由于使用tensorflow总是一种乐趣，因此我们不能直接导出模型，因为标记器包含在模型定义中。不幸的是，核心ONNX平台还不支持这些字符串操作。

幸运的是，ONNXRuntime CustomOps库提供了帮助。这个库也由ONNX团队维护，并为扩展ONNX基本功能的额外定制操作提供支持。你需要安装CMake>3.17.0，才能使用pip install git编译和安装此版本+https://github.com/microsoft/ort-customops.git.

安装Custom Ops库之后，我们将 USE下载到某个文件夹中，并为tf2onnx库提供输出路径。除此之外，我们还可以直接导出模型：

#!/bin/bash
mkdir universal-sentence-encoder-5
cd universal-sentence-encoder-5
wget https://storage.googleapis.com/tfhub-modules/google/universal-sentence-encoder-large/5.tar.gz
tar -xvzf 5.tar.gz
rm 5.tar.gz
cd ..
python -m tf2onnx.convert --saved-model universal-sentence-encoder-5 --output universal-sentence-encoder-5.onnx --opset 12 --extra_opset ai.onnx.contrib:1 --tag serve

tf2onnx库提供了其他一些很好的功能。例如，--signature_def参数允许你部分导出具有多个签名的模型，例如USE v3 For QA。看看这里的参数：https://github.com/onnx/tensorflow-onnx/blob/master/tf2onnx/convert.py

由于底层图和附加的Ops的不同，对USE运行推理现在有些不同。我们必须将Custom Ops库路径传递给ONNX SessionOptions对象。

from onnxruntime import InferenceSession, SessionOptions
from onnxruntime_customops import get_library_path

opt = rt.SessionOptions()
opt.register_custom_ops_library(get_library_path())

sess = rt.InferenceSession("universal-sentence-encoder-5.onnx", opt, providers=ONNX_PROVIDERS)

sess.run(
    output_names=["outputs"],
    input_feed={"inputs:0": [span]},
)[0]

我们的另一个模型

4.试图导出GPT Neo

GPT Neo刚刚在Transformer库发布。它本质上是OpenAI的GPT3架构的开源变体。该模型有两种体系结构：1.3B和2.7B，表示内部参数的数量。模型可通过model-hub获得，https://huggingface.co/EleutherAI。注意：从今天起，你需要transformers-4.5.0.dev0，因为GPT Neo不包含在当前的Pypi包中。

我们首先复制本教程步骤2中的简单导出。这可能有点老套，可能不适合你的环境，因为这一步会丢弃除logits之外的所有输出。但是我们可以从实际硬件上的推理速度方面看到一些数字。

在2021年4月5日，Transformer库提供的完整形状推断似乎没有达到预期的效果，因此我们需要稍作调整。我们只在它周围包装一个自定义层，它返回logits。加载模型需要3分钟的时间，因为我们必须使用外部数据格式来补偿较大的模型大小。再次运行前面的推断：

• TransformerCUDA:114 ms ± 20 ms
• ONNX CUDA（V100）：314 ms ± 4.15 ms
• ONNX TensorRT（V100，ExecutionProvider）：初/workspace/onnxruntime/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc:777 SubGraphCollection_t onnxruntime::TensorrtExecutionProvider::GetSupportedList(SubGraphCollection_t, int, int, const onnxruntime::GraphViewer&, bool*) const [ONNXRuntimeError] : 1 : FAIL : TensorRT input: 649 has no shape specified.

好吧，我相当不满意。但我猜在我们可以正确地导出模型之前，还有更多的优化需要在模型中完成。我不清楚是什么原因导致了这个问题。然而，如果我们查看日志，我们可以看到正在发生的事情:

CUDA kernel not found in registries for Op type: Pad node name: Pad_4368
CUDA kernel not found in registries for Op type: Pad node name: Pad_3801
CUDA kernel not found in registries for Op type: LessOrEqual node name: LessOrEqual_7094
Force fallback to CPU execution for node: Gather_5
Force fallback to CPU execution for node: Unsqueeze_17
Force fallback to CPU execution for node: Slice_37
Force fallback to CPU execution for node: Squeeze_38
Force fallback to CPU execution for node: Div_66
Force fallback to CPU execution for node: Add_35901

这些信息成千上万地出现。我们总共收到10609信息：

这里的关键是:导出到ONNX是一件好事。如果你的模型使用很多当前不支持的Ops，那么它们中的许多都是在CPU上运行的。虽然总的来说这肯定不容易避免，但优化模型就要从它的第一行代码开始。从一开始就要记住你要如何优化它。

结论

在本文中，我们深入研究了ONNX以及如何从pytorch和tensorflow导出模型。现在你可以直接从pytorch自定义和导出模型。你还可以将tensorflow模型从具有自定义操作的检查点导出到ONNX。此外，你还学会了寻找特殊情况。

附加信息

这篇文章的笔记本可以在这里找到：

https://github.com/oborchers/Medium_Repo/blob/master/Putting%20GPT-Neo%20into%20Production%20using%C2%A0ONNX/ONNX-Export.ipynb