• 实验记录

已经安装的环境

1. torch 1.10.1
2. cuda 11.6
3. transformer 3.2.0
4. tensorrt 8.5.2.2

• 安装环境还有一堆坑。

1. pytorch 模型转onnx

# 简单的模型测试, torch模型保存格式为pth
import torch
from transformers import BertTokenizer, BertModel
import torch.nn as nn
import torch.nn.functional as F
import onnx
import onnxruntime


class Model(nn.Module):
    def __init__(self, bert_path, num_classes):
        super(Model, self).__init__()
        self.bert = BertModel.from_pretrained(bert_path)
        self.bert_config = self.bert.config
        out_dims = self.bert_config.hidden_size
        self.fc = nn.Linear(out_dims, num_classes)

    def forward(self, input_ids, token_type_ids, attention_mask):
        _, pooled = self.bert(input_ids, token_type_ids=token_type_ids,
                              attention_mask=attention_mask) 
        out = self.fc(pooled)
        output = F.softmax(out, dim=1)
        return output


def predict_data_process(sentences):
    encoded_input = tokenizer.batch_encode_plus(sentences,
                                                max_length=64,
                                                padding='max_length',
                                                truncation='longest_first')

    input_ids = encoded_input.get('input_ids')
    token_type_ids = encoded_input.get('token_type_ids')
    attention_mask = encoded_input.get('attention_mask')
    input_ids = torch.LongTensor(input_ids).to(DEVICE)
    token_type_ids = torch.LongTensor(token_type_ids).to(DEVICE)
    attention_mask = torch.LongTensor(attention_mask).to(DEVICE)

    return_data = {
        "inputs_ids": input_ids,
        "attention_mask": attention_mask,
        "token_type_ids": token_type_ids,
    }

    return return_data


# 转onnx
def trans2onnx(torch_model, onnx_model_path):

    torch.onnx.export(torch_model,
                      (input_ids, token_type_ids, attention_mask),
                      onnx_model_path,
                      export_params=True,
                      opset_version=11,
                      do_constant_folding=True,
                      # 输入列名
                      input_names=['input_ids', 'token_type_ids', 'attention_mask'],
                      output_names=['output'],
                      # 动态输入，支持batch处理
                      dynamic_axes={'input_ids': {0: 'batch_size', 1: 'seq_len'},
                                    'token_type_ids': {0: 'batch_size', 1: 'seq_len'},
                                    'attention_mask': {0: 'batch_size', 1: 'seq_len'},
                                    'output': {0: 'batch_size', 1: 'num_classes'}})

# onnx 模型验证
def onnx_model_predict(model_path):

    onnx_model = onnx.load(model_path)
    onnx.checker.check_model(onnx_model)

    ort_session = onnxruntime.InferenceSession(model_path)

    # 计算ONNX输出
    ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(input_ids),
                  ort_session.get_inputs()[1].name: to_numpy(token_type_ids),
                  ort_session.get_inputs()[2].name: to_numpy(attention_mask)}
    ort_outs = ort_session.run(None, ort_inputs)
    ort_outs_ = ort_outs[0]
    return ort_outs_


def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()


if __name__ == '__main__':

    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # torch.cuda.set_device(1)
    # 预训练模型路径
    pretrain_model_path = '1'
    # 保存的模型路径
    model_path = "1.pth"
    # 这里10是模型最后一层输出类别数
    torch_model = Model(pretrain_model_path, 10)
    torch_model.load_state_dict(torch.load(model_path, map_location=DEVICE), False)
    torch_model.to(DEVICE)
    tokenizer = BertTokenizer(vocab_file=pretrain_model_path + "/vocab.txt")

    sentences = ['你好', '我好']
    inputs = predict_data_process(sentences)
    input_ids = inputs['inputs_ids']
    token_type_ids = inputs['token_type_ids']
    attention_mask = inputs['attention_mask']

    # pth转onnx
    onnx_model_path = "1.onnx"
    trans2onnx(torch_model, onnx_model_path)
    # pth模型输出
    torch_out = torch_model(input_ids, token_type_ids, attention_mask)
    print("torch_out:", torch_out)
    # 转换后的onnx模型输出
    onnx_out = onnx_model_predict(onnx_model_path)
    print("ort_outs:", onnx_out)
    # 对比torch_out和onnx_out 输出值，就知道模型转换是否有问题。
    

2. onnx 模型转 trt

两种方式：

1. 通过执行trtexec (建议)
   进入安装的Tensor RT目录下的bin文件夹下，比如…/TensorRT-8.5.2.2/bin/ 下执行
   • 注意这里维度是用x号，不是*号。
     trtexec --onnx=your.onnx --saveEngine=your.trt --minShapes=\'input_ids\':1x64,\'token_type_ids\':1x64,\'attention_mask\':1x64 --optShapes=\'input_ids\':32x64,\'token_type_ids\':32x64,\'attention_mask\':32x64 --maxShapes=\'input_ids\':64x64,\'token_type_ids\':64x64,\'attention_mask\':64x64 --fp16 --device=1
     参数：
     –onnx：你的onnx 文件路径
     –saveEngine：你要保存的trt文件路径
     –minShapes：指定输入的最小形状，1 x 1：batch=1，seq_len=1
     –optShapes：指定输入的优化形状，32 x 300：batch=32，seq_len=64
     –maxShapes：指定输入的最大形状，32 x 300：batch=32，seq_len=64
     –fp：转换精度
     –device：指定GPU设备号
2. python 代码

import tensorrt as trt
import os


def convert_onnx_to_engine(onnx_file_path, engine_file_path=None, max_batch_size=32):
    logeer = trt.Logger(trt.Logger.WARNING)
    EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)


    def build_engine():
        with trt.Builder(logeer) as builder, \
                builder.create_network(EXPLICIT_BATCH) as network, \
                builder.create_builder_config() as config, \
                trt.OnnxParser(network, logeer) as parser, \
                trt.Runtime(logeer) as runtime:

            profile = builder.create_optimization_profile()
            # 输入的各列，最小、最优、最大值
            profile.set_shape("input_ids", (1, 64), (32, 64), (32, 64))
            profile.set_shape("token_type_ids", (1, 64), (32, 64), (32, 64))
            profile.set_shape("attention_mask", (1, 64), (32, 64), (32, 64))

            config.add_optimization_profile(profile)

            config.max_workspace_size = 1 << 30
            builder.max_batch_size = max_batch_size
            if not os.path.exists(onnx_file_path):
                print(
                    "ONNX file {} not found, first to generate it.".format(onnx_file_path)
                )
                exit(0)
            print("Loading ONNX file from path {}...".format(onnx_file_path))

            print("Parsing ONNX file")
            with open(onnx_file_path, 'rb') as model:
                if not parser.parse(model.read()):
                    print("ERROR: Failed to parse the ONNX file.")
                    for error in range(parser.num_errors):
                        print(parser.get_error(error))

            print('Completed parsing of ONNX file')

            print("Building TensorRT engine, This may take a few minutes.")
            plan = builder.build_serialized_network(network, config)
            engine = runtime.deserialize_cuda_engine(plan)
            print("Completed creating Engine")
            with open(engine_file_path, "wb") as f:
                f.write(plan)
            return engine

    if os.path.exists(engine_file_path):
        # If a serialized engine exists, use it instead of building an engine.
        print("Reading engine from file {}".format(engine_file_path))
        with open(engine_file_path, "rb") as f, trt.Runtime(logeer) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    else:
        return build_engine()


if __name__ == '__main__':
    onnx_file = "your.onnx"
    engine_file = "your.trt"
    convert_onnx_to_engine(onnx_file, engine_file)

3. trt 模型推理

import torch
import argparse
import numpy as np
import tensorrt as trt
import common
from transformers import BertTokenizer

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)


def get_engine(engine_file_path):
    print("Reading engine from file {}".format(engine_file_path))
    with open(engine_file_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime:
        engine = runtime.deserialize_cuda_engine(f.read())
        return engine


def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()


def predict_data_process(tokenizer, sentences):
    encoded_input = tokenizer.batch_encode_plus(sentences,
                                                max_length=64,
                                                padding='max_length',
                                                truncation='longest_first')

    input_ids = encoded_input.get('input_ids')
    token_type_ids = encoded_input.get('token_type_ids')
    attention_mask = encoded_input.get('attention_mask')
    input_ids = torch.LongTensor(input_ids).to(DEVICE)
    token_type_ids = torch.LongTensor(token_type_ids).to(DEVICE)
    attention_mask = torch.LongTensor(attention_mask).to(DEVICE)

    return_data = {
        "inputs_ids": input_ids,
        "attention_mask": attention_mask,
        "token_type_ids": token_type_ids,
    }

    return return_data


def infer(args, sentences):

    tokenizer = BertTokenizer(vocab_file=args.bert_model + "/vocab.txt")
    inputs = predict_data_process(tokenizer, sentences)
    inputs_ids = to_numpy(inputs['inputs_ids'].int())
    token_type_ids = to_numpy(inputs['token_type_ids'].int())
    attention_mask = to_numpy(inputs['attention_mask'].int())

    """
    a.获取engine, 建立上下文
    """
    engine = get_engine(args.trt_model)
    context = engine.create_execution_context()
    """
    b.从engine中获取inputs, outputs, bindings, stream 的格式以及分配缓存
    """
    context.active_optimization_profile = 0
    origin_inputshape = context.get_binding_shape(0)  # (1,-1)
    origin_inputshape[0], origin_inputshape[1] = inputs_ids.shape  # (batch_size, max_sequence_length)
    # 我模型输入是[inputs_ids, token_type_ids, attention_mask]这三列，如果输入特征列只有两列，去掉一列。
    context.set_binding_shape(0, (origin_inputshape))
    context.set_binding_shape(1, (origin_inputshape))
    context.set_binding_shape(2, (origin_inputshape))

    """
    c、输入数据填充
    """
    inputs, outputs, bindings, stream = common.allocate_buffers(engine, origin_inputshape[0])
    inputs[0].host = inputs_ids
    inputs[1].host = token_type_ids
    inputs[2].host = attention_mask

    """
    d、tensorrt推理
    """
    trt_outputs = common.do_inference_v2(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
    trt_outputs = trt_outputs[0].reshape(origin_inputshape[0], -1)  #
    ort_outs_logits = torch.tensor(trt_outputs)
    _, pred = torch.max(ort_outs_logits.data, 1)
# 这里输出的是模型预测标签的ID
    print("====preds====:\n", pred)

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--trt_model', default='your.trt')
    parser.add_argument('--bert_model', default='your_bert_model')
    parser.add_argument('--output_dir', default='')
    parser.add_argument('--batch_size', default=1, type=int)
    args = parser.parse_args()
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    sentences = ['你好', "我不好"]
    infer(args, sentences)

遇到的问题：

1. 一个警告，不影响后续推理。

Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32. 

网上的查的解决方法： pip install
onnx-simplifier python -m onnxsim your_model.onnx your_model_sim.onnx
我试了，还是有这个警告
2.

pycuda._driver.LogicError: cuMemcpyHtoDAsync failed: invalid argument

当我的模型输入是bathc的时候，会报上述错误，检查发现
comfile.file 里面allocate_buffers()函数下engine.max_batch_size值总是1，所以我把这个函数下engine.max_batch_size 用外部的batch_size值传进来。

附上 common.py

import argparse
import os
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt

try:
    # Sometimes python does not understand FileNotFoundError
    FileNotFoundError
except NameError:
    FileNotFoundError = IOError

EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

def GiB(val):
    return val * 1 << 30


def add_help(description):
    parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    args, _ = parser.parse_known_args()


def find_sample_data(description="Runs a TensorRT Python sample", subfolder="", find_files=[], err_msg=""):
    '''
    Parses sample arguments.
    Args:
        description (str): Description of the sample.
        subfolder (str): The subfolder containing data relevant to this sample
        find_files (str): A list of filenames to find. Each filename will be replaced with an absolute path.
    Returns:
        str: Path of data directory.
    '''

    # Standard command-line arguments for all samples.
    kDEFAULT_DATA_ROOT = os.path.join(os.sep, "usr", "src", "tensorrt", "data")
    parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-d", "--datadir", help="Location of the TensorRT sample data directory, and any additional data directories.", action="append", default=[kDEFAULT_DATA_ROOT])
    args, _ = parser.parse_known_args()

    def get_data_path(data_dir):
        # If the subfolder exists, append it to the path, otherwise use the provided path as-is.
        data_path = os.path.join(data_dir, subfolder)
        if not os.path.exists(data_path):
            if data_dir != kDEFAULT_DATA_ROOT:
                print("WARNING: " + data_path + " does not exist. Trying " + data_dir + " instead.")
            data_path = data_dir
        # Make sure data directory exists.
        if not (os.path.exists(data_path)) and data_dir != kDEFAULT_DATA_ROOT:
            print("WARNING: {:} does not exist. Please provide the correct data path with the -d option.".format(data_path))
        return data_path

    data_paths = [get_data_path(data_dir) for data_dir in args.datadir]
    return data_paths, locate_files(data_paths, find_files, err_msg)

def locate_files(data_paths, filenames, err_msg=""):
    """
    Locates the specified files in the specified data directories.
    If a file exists in multiple data directories, the first directory is used.
    Args:
        data_paths (List[str]): The data directories.
        filename (List[str]): The names of the files to find.
    Returns:
        List[str]: The absolute paths of the files.
    Raises:
        FileNotFoundError if a file could not be located.
    """
    found_files = [None] * len(filenames)
    for data_path in data_paths:
        # Find all requested files.
        for index, (found, filename) in enumerate(zip(found_files, filenames)):
            if not found:
                file_path = os.path.abspath(os.path.join(data_path, filename))
                if os.path.exists(file_path):
                    found_files[index] = file_path

    # Check that all files were found
    for f, filename in zip(found_files, filenames):
        if not f or not os.path.exists(f):
            raise FileNotFoundError("Could not find {:}. Searched in data paths: {:}\n{:}".format(filename, data_paths, err_msg))
    return found_files

# Simple helper data class that's a little nicer to use than a 2-tuple.
class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()

# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
def allocate_buffers(engine, batch_size):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()
    for binding in engine:
        # size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size  # max_batch_size=1
        # print("get_binding_shape", engine.get_binding_shape(binding))
        # print("max_bactch_size", engine.max_batch_size)
        # 这里的修改
        # size = abs(trt.volume(engine.get_binding_shape(binding))) * engine.max_batch_size
        size = abs(trt.volume(engine.get_binding_shape(binding))) * batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # Allocate host and device buffers
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes) # # nbytes表示数组中的所有数据消耗掉的字节数
        # Append the device buffer to device bindings.
        bindings.append(int(device_mem))
        # Append to the appropriate list.
        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))
    return inputs, outputs, bindings, stream


# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
def allocate_buffers_v2(engine, context):
    """
    Allocates host and device buffer for TRT engine inference.
    This function is similiar to the one in ../../common.py, but
    converts network outputs (which are np.float32) appropriately
    before writing them to Python buffer. This is needed, since
    TensorRT plugins doesn't support output type description, and
    in our particular case, we use NMS plugin as network output.
    Args:
        engine (trt.ICudaEngine): TensorRT engine
    Returns:
        inputs [HostDeviceMem]: engine input memory
        outputs [HostDeviceMem]: engine output memory
        bindings [int]: buffer to device bindings
        stream (cuda.Stream): cuda stream for engine inference synchronization
    """
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()
    for i, binding in enumerate(engine):
        # binding:input_ids,input_mask,output
        print(context.get_binding_shape(i)) # (input_ids,input_mask,output).shape (1,105)
        size = trt.volume(context.get_binding_shape(i)) # 1*105
        # dims = engine.get_binding_shape(binding)
        # if dims[1] < 0:
           # size *= -1
        dtype = trt.nptype(engine.get_binding_dtype(binding)) # DataType.FLOAT
        # print(dtype)  # <class 'numpy.float32'>
        # Allocate host and device buffers
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        # Append the device buffer to device bindings.
        bindings.append(int(device_mem))
        # Append to the appropriate list.
        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))
    return inputs, outputs, bindings, stream

# This function is generalized for multiple inputs/outputs.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
    # Transfer input data to the GPU.
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # Run inference. batch_size = 1
    context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
    # Transfer predictions back from the GPU.
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # Synchronize the stream
    stream.synchronize()
    # Return only the host outputs.
    return [out.host for out in outputs]

# This function is generalized for multiple inputs/outputs for full dimension networks.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference_v2(context, bindings, inputs, outputs, stream):
    # Transfer input data to the GPU.
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # Run inference.
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    # Transfer predictions back from the GPU.
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # Synchronize the stream
    stream.synchronize()
    # Return only the host outputs.
    return [out.host for out in outputs]