自己写的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_