onnx格式转tensorRT

自己写的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_
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值