导出BERT句子模型为ONNX并推理

在深度学习中，将模型导出为ONNX（Open Neural Network Exchange）格式并利用ONNX进行推理是提高推理速度和模型兼容性的一种常见做法。本文将介绍如何将BERT句子模型导出为ONNX格式，并使用ONNX Runtime进行推理，具体以中文文本处理为例。

1. 什么是ONNX？

ONNX 是一种开放的神经网络交换格式，旨在促进深度学习模型在不同平台和工具之间的共享和移植。它支持包括PyTorch、TensorFlow等多种主流框架，可以通过ONNX Runtime库高效推理。通过将模型转换为ONNX格式，我们可以获得跨平台部署的优势，并利用ONNX Runtime加速推理过程。

2. 准备工作

在导出和推理之前，需要安装以下库：

pip install torch transformers onnx onnxruntime

3. 导出BERT句子模型为ONNX

首先，我们将使用HuggingFace的transformers库加载一个预训练的BERT句子模型（text2vec-base-chinese），然后将其导出为ONNX格式。以下是导出模型的步骤和代码：

3.1 导出模型的代码

import torch
from transformers import BertTokenizer, BertModel

# 加载预训练的BERT模型和分词器
tokenizer = BertTokenizer.from_pretrained('shibing624/text2vec-base-chinese')
model = BertModel.from_pretrained('shibing624/text2vec-base-chinese')

# 读取要处理的句子
with open("corpus/words_nlu.txt", 'rt', encoding='utf-8') as f:
    nlu_words = [line.strip() for line in f.readlines()]
nlu_words.insert(0, "摄像头打开一下")  # 插入要比较的句子

# 对句子进行编码
encoded_input = tokenizer(nlu_words, padding=True, truncation=True, return_tensors='pt')

# 设置ONNX模型的保存路径
onnx_model_path = "text2vec-base-chinese.onnx"
model.eval()

# 导出模型为ONNX格式
with torch.no_grad():
    torch.onnx.export(
        model,
        (encoded_input['input_ids'], encoded_input['attention_mask']),
        onnx_model_path,
        input_names=['input_ids', 'attention_mask'],
        output_names=['last_hidden_state'],
        opset_version=14,
        dynamic_axes={
            'input_ids': {0: 'batch_size', 1: 'sequence_length'},
            'attention_mask': {0: 'batch_size', 1: 'sequence_length'},
            'last_hidden_state': {0: 'batch_size', 1: 'sequence_length'}
        }
    )
print(f"ONNX模型已导出到 {onnx_model_path}")

在这段代码中，我们将text2vec-base-chinese模型导出为ONNX格式，指定了输入和输出的名称，并使用了动态轴设置（如批大小和序列长度），这样可以处理不同长度的句子。

4. 使用ONNX进行推理

导出模型后，我们可以使用ONNX Runtime进行推理。以下是基于ONNX的推理代码。该代码实现了对输入文本进行预处理、调用ONNX模型进行推理、以及对模型输出进行均值池化处理。

4.1 ONNX推理代码

import os
from onnxruntime import InferenceSession
import numpy as np

class Text2Vector: 
    def __init__(self, model_dir="text2vec-base-chinese") -> None:
        self.model_path = os.path.join(model_dir,"text2vec-base-chinese.onnx")
        self.vocab_path = os.path.join(model_dir,"vocab.txt")
        self.vocab = self.load_vocab(self.vocab_path)
        self.onnx_session = InferenceSession(self.model_path)
        

    def load_vocab(self, vocab_path):
        """Load BERT vocabulary."""
        vocab = {}
        with open(vocab_path, 'r', encoding='utf-8') as f:
            for idx, line in enumerate(f):
                token = line.strip()
                vocab[token] = idx
        return vocab

    def tokenize(self, text):
        """Tokenize text into BERT input_ids."""
        tokens = ['[CLS]']
        for char in text:
            if char in self.vocab:
                tokens.append(char)
            else:
                tokens.append('[UNK]')
        tokens.append('[SEP]')
        input_ids = [self.vocab[token] if token in self.vocab else self.vocab['[UNK]'] for token in tokens]
        return input_ids

    def preprocess(self, texts, max_length=128):
        """Preprocess input texts for BERT model."""
        input_ids_list = []
        attention_mask_list = []
        
        for text in texts:
            input_ids = self.tokenize(text)
            # Truncate or pad to max_length
            if len(input_ids) > max_length:
                input_ids = input_ids[:max_length]
            else:
                input_ids += [0] * (max_length - len(input_ids))

            attention_mask = [1 if idx != 0 else 0 for idx in input_ids]
            
            input_ids_list.append(input_ids)
            attention_mask_list.append(attention_mask)

        # Convert to NumPy arrays
        inputs = {
            'input_ids': np.array(input_ids_list, dtype=np.int64),
            'attention_mask': np.array(attention_mask_list, dtype=np.int64)
        }
        return inputs

    def mean_pooling_numpy(self, model_output, attention_mask):
        """Mean pooling for model output."""
        token_embeddings = model_output
        input_mask_expanded = np.expand_dims(attention_mask, -1).astype(float)
        return np.sum(token_embeddings * input_mask_expanded, axis=1) / np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)

    def compute_embeddings(self, texts):
        """Compute sentence embeddings for input texts."""
        onnx_inputs = self.preprocess(texts)
        # Run the model with ONNX Runtime
        onnx_outputs = self.onnx_session.run(None, onnx_inputs)
        last_hidden_state = onnx_outputs[0]
        # Perform mean pooling
        sentence_embeddings = self.mean_pooling_numpy(last_hidden_state, onnx_inputs['attention_mask'])
        return sentence_embeddings

    def match(self, cmd_embeddings, text_embedding):

        # Compute cosine similarity between each command embedding and the text embedding
        cmd_embeddings_norm = cmd_embeddings / np.linalg.norm(cmd_embeddings, axis=1, keepdims=True)
        text_embedding_norm = text_embedding / np.linalg.norm(text_embedding, axis=1,keepdims=True)
        
        # Cosine similarity calculation
        similarity_scores = np.dot(cmd_embeddings_norm, text_embedding_norm.T).squeeze()
        # Find the best match index and score
        best_match_idx = np.argmax(similarity_scores)
        best_score = similarity_scores[best_match_idx]
        
        return best_match_idx, best_score

def main():
    # Initialize the PIPE_NLU instance
    nlu = Text2Vector(model_dir=r"D:\code\text2vec-base-chinese-onnx")

    # Define a list of command sentences
    commands = [
        "打开音乐",
        "播放视频",
        "关闭灯光",
        "查询天气",
        "设置闹钟"
    ]

    # Define a target sentence to match against the command sentences
    target_sentence = "请播放音乐"

    # Compute embeddings for commands and the target sentence
    cmd_embeddings = nlu.compute_embeddings(commands)
    target_embedding = nlu.compute_embeddings([target_sentence])

    # Match the target sentence against the command embeddings
    best_match_idx, best_score = nlu.match(cmd_embeddings, target_embedding)

    # Output the results
    print(f"Best match command: {commands[best_match_idx]}")
    print(f"Similarity score: {best_score:.4f}")

if __name__ == "__main__":
    main()

4.2 推理流程

1. 加载ONNX模型：通过InferenceSession加载ONNX模型。
2. 加载词汇表：读取BERT的词汇表，用于将输入文本转化为模型可接受的input_ids格式。
3. 文本预处理：将输入的文本进行分词、截断或填充为固定长度，并生成相应的注意力掩码attention_mask。
4. 模型推理：通过ONNX Runtime调用模型，获取句子的最后隐藏状态输出。
5. 均值池化：对最后的隐藏状态进行均值池化，计算出句子的嵌入向量。
6. 归一化嵌入：将句子嵌入向量进行归一化，使得向量长度为1。

5. 总结

通过将BERT模型导出为ONNX并使用ONNX Runtime进行推理，我们可以大幅度提升推理速度，同时保持了高精度的句子嵌入计算。在实际应用中，ONNX Runtime的跨平台特性和高性能表现使其成为模型部署和推理的理想选择。

使用上述步骤，您可以轻松将BERT句子模型应用到各种自然语言处理任务中，如语义相似度计算、文本分类和句子嵌入等。