pytorch构建的深度学习模型（pt文件）转换为onnx格式，并支持batch输入，以bert模型为例

在研发深度学习模型过程中，通过pytorch、tensorflow、paddlepaddle等深度学习框架构建深度学习模型，但在模型部署时希望模型格式和部署环境能够相对统一。因此，考虑将深度学习模型转换为onnx格式。通过onnxruntime进行模型推理，或者再转换为tensorRT、paddlepaddle等格式进一步加速。NLP预训练模型出来后，bert类模型推理时速度慢。huggingface的transformers底层基于pytorch，因此记录一下pytorch构建的bert转onnx

---

1. 基本步骤

（1）考虑到模型的训练流程和推理（测试）流程可能有所不同，先载入训练好的网络模型参数，执行模型推理过程。

（2）安装onnx

（3）执行torch.onnx.export方法生成onnx文件。如果想要batch推理且模型输入有多个变量（如bert 模型有input_ids, attention_masks, token_type_ids）需要进行动态轴的设置。

2. 版本依赖

        本文采用的版本为cuda10.1，pytorch 1.7.1，transformers 2.10.0，onnx 1.8.0 , onnxruntime 1.4.0，pytorch 1.7.1的torch.onnx方法可见以下链接： torch.onnx — PyTorch 1.7.1 documentation

        还需要安装onnx 和 onnxruntime：

pip install onnx == 1.8
#pip install onnx

        只能转换torch.onnx版本支持的算子，如果构建的网络结构的底层算子，不在torch.onnx的支持里面就没法正常转换了，需要自己构建自定义算子，但大多数使用者没有这样的开发能力。随着版本的不断升级完善许多算子后续就支持了，也不需要花太多时间，最好使用支持的算子来构建网络。

        onnx支持的算子：https://github.com/onnx/onnx/blob/master/docs/Operators.md

不同的onnx opset版本集合支持的算子不一样，需要将版本设置成支持模型所有算子操作的版本

---

3. 模型转换

    3.1 导入环境

import torch
import os
from transformers import BertTokenizer
import onnxruntime as ort
import numpy as np

    3.2 模型推理

GPU_IDS = '-1' #模型格式转换可以不用GPU

model, device = load_model_and_parallel(model, GPU_IDS, CKPT_PATH)
#CKPT_PATH是训练好的bert模型pt文件的地址
model.eval()
tokenizer = BertTokenizer('vocab.txt') #bert词表
with torch.no_grad():
    sen = '今天天气很好'
    sent_tokens = fine_grade_tokenize(sen , tokenizer)
    encode_dict = tokenizer.encode_plus(text=sent_tokens,
                                        max_length=MAX_SEQ_LEN,
                                        is_pretokenized=True,
                                        pad_to_max_length=True,
                                        return_tensors='pt',
                                        return_token_type_ids=True,
                                        return_attention_mask=True)
    samples = {'token_ids': encode_dict['input_ids'],
               'attention_masks': encode_dict['attention_masks'],
               'token_type_ids': encode_dict['token_type_ids']}

    for sample in samples:
        samples[sample] =  samples[sample].to(device)
    
    output1,output2  = model(**samples) #模型输出可能有多个,此处假设有两个

fine_grade_tokenize和load_model_and_parallel函数，参考于以下github。fine_grade_tokenize函数，如果需要进行序列标注相关任务可以参考，仅仅分类任务可以不做这样的处理：https://github.com/RongkaiPeng/DeepNER/blob/master/src/preprocess/processor.py

def fine_grade_tokenize(raw_text, tokenizer):
    """
    序列标注任务 BERT 分词器可能会导致标注偏移，
    用 char-level 来 tokenize
    """
    tokens = []

    for _ch in raw_text:
        if _ch in [' ', '\t', '\n']:
            tokens.append('[BLANK]')
        else:
            if not len(tokenizer.tokenize(_ch)):
                tokens.append('[INV]')
            else:
                tokens.append(_ch)

    return tokens

def load_model_and_parallel(model, gpu_ids, ckpt_path=None, strict=True):
    """
    加载模型 & 放置到 GPU 中（单卡 / 多卡）
    """
    gpu_ids = gpu_ids.split(',')

    # set to device to the first cuda
    device = torch.device("cpu" if gpu_ids[0] == '-1' else "cuda:" + gpu_ids[0])

    if ckpt_path is not None:
        logger.info(f'Load ckpt from {ckpt_path}')
        model.load_state_dict(torch.load(ckpt_path, map_location=torch.device('cpu')), strict=strict)

    model.to(device)

    if len(gpu_ids) > 1:
        logger.info(f'Use multi gpus in: {gpu_ids}')
        gpu_ids = [int(x) for x in gpu_ids]
        model = torch.nn.DataParallel(model, device_ids=gpu_ids)
    else:
        logger.info(f'Use single gpu in: {gpu_ids}')

    return model, device

    3.3  转换为onnx文件

          通过torch.repeat_interleave，模拟batch格式的输入，转换成Tensor

sen = '今天天气很好'
sent_tokens = fine_grade_tokenize(sen , tokenizer)
encode_dict = tokenizer.encode_plus(text=sent_tokens,
                                    max_length=MAX_SEQ_LEN,
                                    is_pretokenized=True,
                                    pad_to_max_length=True,
                                    return_tensors='pt',
                                    return_token_type_ids=True,
                                    return_attention_mask=True)

batch_size = 32

input1 = encode_dict['input_ids']
input2 = encode_dict['attention_masks']
input3 = encode_dict['token_type_ids']
batch_input1 = torch.repeat_interleave(input1, repeats = batch_size, dim = 0)
batch_input2 = torch.repeat_interleave(input2, repeats = batch_size, dim = 0)
batch_input3 = torch.repeat_interleave(input3, repeats = batch_size, dim = 0)

tensor_input0 = torch.LongTensor(batch_input1)
tensor_input1 = torch.LongTensor(batch_input2)
tensor_input2 = torch.LongTensor(batch_input3)

        通过torch.onnx.export接口进行转换，其中input_names,out_names是模型的输入输出在模型构建时取的名字。

input_names = ['input_ids', 'input_masks', 'token_type_ids']
outputs_names = ['output1', 'output2']
onnx_name = 'convert.onnx'
torch.onnx.export(model, (), onnx_name, input_names = input_names , 
                  output_names = outputs_names , verbose = True, opset_version = 11,
                  dynamic_axes = {'input_ids':{0: 'batch_size'},
                                 {'input_masks':{0: 'batch_size'},
                                 {'token_type_ids':{0: 'batch_size'},
                                 {'output1':{0: 'batch_size'},
                                 {'output2':{0: 'batch_size'}})

          batch输入必须设置dynamic_axes，不然batch推理就不行了