Kerberos认证原理 leetcodeLCP wxRuby 反射 Java Spring SCI static stream insert swift2 Avalon vue响应式 纯html网页模板 oracle删除字段sql excel动态图表制作 js获取body的高度 字符串中包含某个字符串 java手机验证码 office配置进度 matlab不等于怎么表示 linux启动数据库 docker启动命令 python迭代器 python输出中文 python基础教程免费 python平台 python的安装路径 python服务器开发 javaqueue java程序实例 java获取当前年份 java语言 java初级入门教程 java数组添加值 怎么安装java环境 java如何配置环境变量 java日期格式化 java异常 修改mac地址软件 网络适配器驱动
当前位置: 首页 > 学习教程  > 编程语言

onnx格式转tensorRT

2021/1/13 19:17:18 文章标签: 测试文章如有侵权请发送至邮箱809451989@qq.com投诉后文章立即删除

自己写的onnx转trt的代码。 此处记录几点: 模型的输入须为numpy格式,所以直接从DataLoader取出的数据是不能直接扔进模型的模型的输入是多个的时候,如输入多张图片时,可以通过下面这种代码格式传入数据。 inputs[0].host data[0…

自己写的onnx转trt的代码。
此处记录几点:

  • 模型的输入须为numpy格式,所以直接从DataLoader取出的数据是不能直接扔进模型的
  • 模型的输入是多个的时候,如输入多张图片时,可以通过下面这种代码格式传入数据。
	inputs[0].host = data[0]
    inputs[1].host = data[1]
    inputs[2].host = data[2]
    inputs[3].host = data[3]

【该代码是在已经获得engine文件和onnx文件的情况下进行使用的。】

此处给出快速获得engine文件的方法:使用TensorRT自带的trtexec,见:
https://zhuanlan.zhihu.com/p/158199822

# --*-- coding:utf-8 --*--
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
import torch
import time
from PIL import Image
import cv2, os
import torchvision
import numpy as np
import model1
from torch.utils.data import Dataset
import torch
import h5py
import random
from torch.utils.data import DataLoader
from torch.autograd import Variable


class NTUData(Dataset):
    def __init__(self, cv_path, transform=None):
        # TODO
        # 1. Initialize file path or list of file names.
        self.cv_path = cv_path
        self.transform = transform
        with h5py.File(self.cv_path, 'r') as file:
            self.group_name_list = [name for name in file]  # Êý¾Ý¼¯h5µÄdata name
            # shape[1]È¡¹Ç¼ÜÐòÁеÄÖ¡Êý
            self.num_per_frame = None
            self.data_len = file.keys().__len__()  # ¹Ç¼ÜÐòÁÐÊý¾ÝµÄÊýÄ¿

    def __getitem__(self, index):
        # TODO
        # 1. Read one data from file (e.g. using numpy.fromfile, PIL.Image.open).
        # 2. Preprocess the data (e.g. torchvision.Transform).
        # 3. Return a data pair (e.g. image and label).
        # ÕâÀïÐèҪעÒâµÄÊÇ£¬µÚÒ»²½£ºread one data£¬ÊÇÒ»¸ödata
        index_30 = self.get_frame30_index_list(index)
        # print(index_30.shape)
        frame30_body_0, frame30_body_1, label_action = self.get_frame30_data(index, index_30)  # get segment data ,label
        diff_body_0, diff_body_1 = self.get_diff_data(frame30_body_0, frame30_body_1)  # get skeleton Motion

        frame30_body_0 = Variable(torch.FloatTensor(frame30_body_0))
        frame30_body_1 = Variable(torch.FloatTensor(frame30_body_1))
        diff_body_0 = Variable(torch.FloatTensor(diff_body_0))
        diff_body_1 = Variable(torch.FloatTensor(diff_body_1))
        label_action = torch.LongTensor(label_action).squeeze().numpy().tolist().index(1)

        return frame30_body_0, frame30_body_1, diff_body_0, diff_body_1, label_action

    def __len__(self):
        # You should change 0 to the total size of your dataset.
        return self.data_len

    def get_frame30_index_list(self, index):
        with h5py.File(self.cv_path, 'r') as file:
            self.num_per_frame = file[self.group_name_list[index]].shape[1]  # ÿ¸ö¹Ç¼ÜÐòÁеÄ×ÜÖ¡Êý
        # Éú³É31¸öÊýÊÇΪÁ˵õ½30¸öÇø¼ä
        random_list = np.linspace(0, self.num_per_frame, 31, dtype=int)  # ºóÃæ»á¼õÒ»£¬²»»áÈ¡µ½µÚ109Ö¡
        # ¼õÒ»ÊÇȡÿһСÇø¼äµÄ×îºóһλÊý£¬±ÜÃâÈ¡µ½ÏÂÒ»Çø¼äµÄµÚÒ»¸öÊý£¬±ÜÃâ³öÏÖÏàͬ֡µÄÇé¿ö
        frame30_list = [random.randint(random_list[i], random_list[i + 1] - 1) for i in range(30)]
        # frame30_list = sorted(random.sample(range(0, self.num_per_frame), 30))  # index list(1, 30)
        # frame30_index_list = sorted(np.arange(0, self.num_per_frame, 3))
        return frame30_list

    def get_frame30_data(self, index, frame30_list):
        with h5py.File(self.cv_path, 'r') as file:
            frame30_body_0 = file[self.group_name_list[index]][0, frame30_list, :, :]  # дÈë30Ö¡Êý¾Ý(30, 25, 3)
            frame30_body_1 = file[self.group_name_list[index]][1, frame30_list, :, :]  # дÈë30Ö¡Êý¾Ý
            label_action = file[self.group_name_list[index]].attrs['label']  # дÈëlabel(60, 1)
        return frame30_body_0, frame30_body_1, label_action

    def get_diff_data(self, body0, body1):
        diff_body_0 = np.diff(body0, n=1, axis=0)  # (29, 25, 3)
        diff_body_1 = np.diff(body1, n=1, axis=0)
        diff_zero = np.zeros((1, 25, 3))
        diff_body_0 = np.r_[diff_body_0, diff_zero]
        diff_body_1 = np.r_[diff_body_1, diff_zero]
        return diff_body_0, diff_body_1


cv_tst_path = '../cv_tst.hdf5'
test_data = NTUData(cv_tst_path)
net = model1.ConVNet().cuda()
net.eval()

# filename = '/home/zp1/code/tensorrt/1.png'
max_batch_size = 1
onnx_model_path = "./ConVNet.onnx"

TRT_LOGGER = trt.Logger()


class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        """
        host_mem: cpu memory
        device_mem: gpu memory
        """
        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__()


def allocate_buffers(engine):
    inputs, outputs, bindings = [], [], []
    stream = cuda.Stream()
    for binding in engine:
        print('the bindding\'s name:', binding)  # 绑定的输入输出
        print('bingding shape: ', engine.get_binding_shape(binding))  # get_binding_shape 是变量的大小
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        # volume 计算可迭代变量的空间,指元素个数
        # size = trt.volume(engine.get_binding_shape(binding)) # 如果采用固定bs的onnx,则采用该句
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # get_binding_dtype  获得binding的数据类型
        # nptype等价于numpy中的dtype,即数据类型
        # allocate host and device buffers
        host_mem = cuda.pagelocked_empty(size, dtype)  # 创建锁业内存
        device_mem = cuda.mem_alloc(host_mem.nbytes)  # cuda分配空间
        # print(int(device_mem)) # binding在计算图中的缓冲地址
        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


def get_engine(max_batch_size=1, onnx_file_path="", engine_file_path="", fp16_mode=False, save_engine=False):
    """
    params max_batch_size:      预先指定大小好分配显存
    params onnx_file_path:      onnx文件路径
    params engine_file_path:    待保存的序列化的引擎文件路径
    params fp16_mode:           是否采用FP16
    params save_engine:         是否保存引擎
    returns:                    ICudaEngine
    """
    # 如果已经存在序列化之后的引擎,则直接反序列化得到cudaEngine
    if os.path.exists(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:
            return runtime.deserialize_cuda_engine(f.read())  # 反序列化
    # else:  # 由onnx创建cudaEngine
    #
    #     # 使用logger创建一个builder
    #     # builder创建一个计算图 INetworkDefinition
    #     # In TensorRT 7.0, the ONNX parser only supports full-dimensions mode, meaning that your network definition must be created with the explicitBatch flag set. For more information, see Working With Dynamic Shapes.
    #     explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    #
    #     with trt.Builder(TRT_LOGGER) as builder, \
    #             builder.create_network(explicit_batch) as network, \
    #             trt.OnnxParser(network, TRT_LOGGER) as parser:  # 使用onnx的解析器绑定计算图,后续将通过解析填充计算图
    #         builder.max_workspace_size = 1 << 30  # 预先分配的工作空间大小,即ICudaEngine执行时GPU最大需要的空间
    #         builder.max_batch_size = max_batch_size  # 执行时最大可以使用的batchsize
    #         builder.fp16_mode = fp16_mode
    #
    #         # 解析onnx文件,填充计算图
    #         if not os.path.exists(onnx_file_path):
    #             quit("ONNX file {} not found!".format(onnx_file_path))
    #         print('loading onnx file from path {} ...'.format(onnx_file_path))
    #         with open(onnx_file_path, 'rb') as model:  # 二值化的网络结果和参数
    #             print("Begining onnx file parsing")
    #             parser.parse(model.read())  # 解析onnx文件
    #         # parser.parse_from_file(onnx_file_path) # parser还有一个从文件解析onnx的方法
    #
    #         print("Completed parsing of onnx file")
    #         # 填充计算图完成后,则使用builder从计算图中创建CudaEngine
    #         print("Building an engine from file{}' this may take a while...".format(onnx_file_path))
    #
    #         #################
    #         print(network.get_layer(network.num_layers - 1).get_output(0).shape)
    #         # network.mark_output(network.get_layer(network.num_layers -1).get_output(0))
    #         last_layer = network.get_layer(network.num_layers - 1)
    #         network.mark_output(last_layer.get_output(0))
    #         engine = builder.build_cuda_engine(network)  # 注意,这里的network是INetworkDefinition类型,即填充后的计算图
    #         print("Completed creating Engine")
    #         if save_engine:  # 保存engine供以后直接反序列化使用
    #             with open(engine_file_path, 'wb') as f:
    #                 f.write(engine.serialize())  # 序列化
    #         return engine


def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
    # Transfer data from CPU to the GPU.
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    # htod: host to device 将数据由cpu复制到gpu device
    # Run inference.
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    # 当创建network时显式指定了batchsize, 则使用execute_async_v2, 否则使用execute_async
    # Transfer predictions back from the GPU.
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # gpu to cpu
    # Synchronize the stream
    stream.synchronize()
    # Return only the host outputs.
    return [out.host for out in outputs]


def postprocess_the_outputs(h_outputs, shape_of_output):
    h_outputs = h_outputs.reshape(*shape_of_output)
    return h_outputs


Lr = 0.0002
Epochs = 300
Batch_Size = 64
test_loader = DataLoader(dataset=test_data, batch_size=Batch_Size, shuffle=False)

# img_np_nchw = get_img_np_nchw(filename).astype(np.float32)
# These two modes are depend on hardwares
fp16_mode = False
trt_engine_path = "./ConVNet.engine"
# Build an cudaEngine
engine = get_engine(max_batch_size, onnx_model_path, trt_engine_path, fp16_mode)
# 创建CudaEngine之后,需要将该引擎应用到不同的卡上配置执行环境
context = engine.create_execution_context()
inputs, outputs, bindings, stream = allocate_buffers(engine)  # input, output: host # bindings

shape_of_output = (64, 4096)

#t_row和t_rt用于计算pytorch和trt的推断时间
t_row = 0
t_rt = 0

# 数据的输入需要为numpy格式,为了更好的测试trt的速度,此处先将所有的数据取出
l = []
for data in test_loader:
    inputs[0].host = data[0].numpy()
    inputs[1].host = data[1].numpy()
    inputs[2].host = data[2].numpy()
    inputs[3].host = data[3].numpy()
    l.append([data[0].numpy(),data[1].numpy(),data[2].numpy(),data[3].numpy()])
    # inputs[1].host = ... for multiple input
t1 = time.time()
for data in l:
    inputs[0].host = data[0]
    inputs[1].host = data[1]
    inputs[2].host = data[2]
    inputs[3].host = data[3]
    trt_outputs = do_inference(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)  # numpy data
    # feat = postprocess_the_outputs(trt_outputs[0], shape_of_output)
t2 = time.time()
t_rt += (t2 - t1)
t3 = time.time()
for data in test_loader:
    output = net(data[0].cuda(), data[1].cuda(), data[2].cuda(), data[3].cuda())
t4 = time.time()
t_row += t4 - t3
print('inference in pytorch 总耗时', t_row)
print('inference in trt 总耗时', t_rt)
print('TensorRT ok')

# mse = np.mean((feat - feat_2) ** 2)
# print('MSE Error = {}'.format(mse))

print('All completed!')


本文链接: http://www.dtmao.cc/news_show_600004.shtml

附件下载

相关教程

    暂无相关的数据...

共有条评论 网友评论

验证码: 看不清楚?