在Python接口中还有一个高度抽象的推理模块叫做TensorRT Lite。TensorRT Lite几乎可以完成任何任务,从创建一个engine,执行推理。因此,用户可以简单的创建一个engine,开始处理数据。
TensorRT Lite接口位于tensorrt.lite,包含一个叫做Engine的类。engine构造利用模型定义,输入输出,从而生成一个完整的可以用来推理运算的engine。
在Lite engine内部自动创建了logger,TensorRTengine,runtime与context,并且分配了GPU显存。
用户自定义的loggers,plugins,calibrators与profilers可以从构造器里传递过来。
不想常用的TensorRT函数,LiteEngine直接支持TensorFlow模型,在内部进行了UFF格式的转换。模型可以以关键字作为参数,序列图或者protobuf文件路径的形式提供。例如:
mnist_engine =tensorrt.lite.Engine(framework="tf", #Source framework
path="mnist/lenet5_mnist_frozen.pb",
#Model File
max_batch_size=10, #Max number of images
to be processed at a time
input_nodes={"in":(1,28,28)}, #Input layers
output_nodes=["out"]) #Output layers
与通过TensorFlow模型创建lite engine类似,Lite Engine接口也可以输入UFF模型stream或者UFF文件。例如:
stream =uff.from_tensorflow_frozen_model(“resnet50-infer-5.pb" ,["GPU_0/
tower_0/Softmax"])
resnet =tensorrt.lite.Engine(framework="uff",
stream=stream,input_nodes={"in":(3,224,224)},output_nodes=["GPU_0/tower_0/
Softmax"])
对于使用NVCaffe模型,用户需要提供deploy文件与模型文件的路径,例如:
c = tensorrt.lite.Engine(framework="caffe",
deployfile= "/mnist/mnist.prototxt",
modelfile= "/mnist/mnist.caffemodel",
max_batch_size=10,
input_nodes={"data":(1,28,28)},
output_nodes=["prob"])
确保你已经有一个TensorRTengine与相应PLAN文件。两者具备有,用户可以使用Lite接口进行导入。例如:
engine =tensorrt.lite.Engine(PLAN="model.plan")
如果用户已经有engine在内存中,或者你手动使用网络定义接口转换了模型,用户可以Lite接口将engine导入:
import tensorrt as trt
import torch
import numpy as np
GLOGGER =trt.infer.ColorLogger(trt.infer.LogSeverity.ERROR)
#Create an Engine from a PyTorch model
def create_pytorch_engine(max_batch_size, builder, dt,model):
network = builder.create_network()
data = network.add_input("in", dt, (1, 28,28))
assert(data)
#-------------
conv1_w =model['conv1.weight'].cpu().numpy().reshape(-1)
conv1_b =model['conv1.bias'].cpu().numpy().reshape(-1)
conv1 = network.add_convolution(data, 20, (5,5),conv1_w, conv1_b)
assert(conv1)
conv1.set_stride((1,1))
#-------------
pool1 = network.add_pooling(conv1.get_output(0),trt.infer.PoolingType.MAX,
(2,2))
assert(pool1)
pool1.set_stride((2,2))
#-------------
conv2_w =model['conv2.weight'].cpu().numpy().reshape(-1)
conv2_b =model['conv2.bias'].cpu().numpy().reshape(-1)
conv2 = network.add_convolution(pool1.get_output(0),50, (5,5), conv2_w,
conv2_b)
assert(conv2)
conv2.set_stride((1,1))
#-------------
pool2 = network.add_pooling(conv2.get_output(0),trt.infer.PoolingType.MAX,
(2,2))
assert(pool2)
pool2.set_stride((2,2))
#-------------
fc1_w = model['fc1.weight'].cpu().numpy().reshape(-1)
fc1_b = model['fc1.bias'].cpu().numpy().reshape(-1)
fc1 = network.add_fully_connected(pool2.get_output(0),500, fc1_w, fc1_b)
assert(fc1)
#-------------
relu1 = network.add_activation(fc1.get_output(0),
trt.infer.ActivationType.RELU)
assert(relu1)
#-------------
fc2_w = model['fc2.weight'].cpu().numpy().reshape(-1)
fc2_b = model['fc2.bias'].cpu().numpy().reshape(-1)
fc2 = network.add_fully_connected(relu1.get_output(0),10, fc2_w, fc2_b)
assert(fc2)
#-------------
# NOTE: Before release
# Using log_softmax in training, cutting out logsoftmax here since no log
softmax in TRT
fc2.get_output(0).set_name("out")
network.mark_output(fc2.get_output(0))
builder.set_max_batch_size(max_batch_size)
builder.set_max_workspace_size(1 << 20)
engine = builder.build_cuda_engine(network)
network.destroy()
return engine
def main():
#Load pretrained PyTorch model
model = torch.load(DATA_DIR +"/mnist/trained_lenet5_mnist.pyt")
#Build an engine from PyTorch
builder = trt.infer.create_infer_builder(GLOGGER)
engine = create_pytorch_engine(10, builder,trt.infer.DataType.FLOAT, model)
#Create a Lite Engine from the engine
mnist_engine =trt.lite.Engine(engine_stream=engine.serialize(), #Use a
serialized engine
max_batch_size=10)
#Max batch size
#Destroy the old engine
engine.destroy()
...
engine创建以后,用户可以使用infer的调用,在数据集上运行推理运算。用户数据有一系列方法传到Lite engine里。总的来说,每个输入都必须是满足构造函数输入尺寸的NumPy矩阵。例如,假设输入层的尺寸是1,28,28,之后每个输入尺寸都必须是1,28,28。从基本结构来说,输入数据可以归为:
Ø 单个输入,例如,一张图片,这时一个三维的NumPy矩阵可以满足输入尺寸。
Ø 一个列表或者NumPy矩阵数据,例如,一系列可以用三维NumPy矩阵表示的图片,每个三维NumPy矩阵可以匹配输入尺寸。
注:图片列表的长度可以随意设置,在内部会根据最大batchsize自动进行分组。
Ø 还有一种情况是batch图片的列表,其中每个batch都必须小于最大batchsize。
如果用户有多个输入层,将不同层的输入作为独立的参数以用户在构造器中定义的顺序传到engine中。每层的格式必须一致,小于batchsize。
推理运行准备好后,结果会以输入格式一致的格式返回。数据将会以输入list相同的索引返回。
更多信息请见Creating A TensorRT Lite Engine。
通常在进行大规模应用的时候,都需要对推理前输入数据进行数据预处理,推理后数据进行后处理。为了应用代码的简洁性,构造器允许用户构建一个用于数据预处理与后处理的函数列表。
例如,如果用户输入大量的原始数据到推理函数中,每张图片都会在进行推理前进行正则化操作。对输出结果进行ArgMax运算得到Top1与Top5的结果。
下列的实例代码,用户创建了包含功能函数的字典,以各自的层名作为关键字。如果用户有多个输入层,但是仅需要预处理一个,用户还是必须对所有输入进行处理,但是对不需要待处理的传入None。
# Preprocessing function
def normalize(data):
#each image is provided as a 3D numpy array (like howit’s provided to
inference function)
for i in range(len(data)): #normalize
data[i] = 1.0 - data[i] / 255.0
#Reshape the data to the shape expected by the network
return data.reshape(1,28,28)
#Lamba to apply argmax to each result after inferenceto get prediction
#Instead of having to reshape, you can replace the 3Darray provided to the
postprocessor with #the object of your choosing (e.g.the top class)
argmax = lambda res: np.argmax(res.reshape(10))
#Register pre and post processors to their layers
mnist_engine =tensorrt.lite.Engine(framework="tf", #Source framework
path=DATA_DIR + "/mnist/
lenet5_mnist_frozen.pb", #Model File
max_batch_size=10, #Max number of images
to be processed at a time
input_nodes={"in":(1,28,28)}, #Input layers
output_nodes=["out"], #Ouput layers
preprocessors={"in":normalize},
#Preprocessing functions
postprocessors={"out":argmax})
#Postprocesssing functions
def generate_cases(num):
'''
Generate a list of raw data (data will be processed inthe engine) and
answers to compare to
'''
cases = []
labels = []
for c in range(num):
rand_file = randint(0, 9)
im = Image.open(str(rand_file) + ".pgm")
arr = np.array(im).reshape(1,28,28) #Make the imageCHANNEL x HEIGHT x
WIDTH
cases.append(arr) #Append the image to list of imagesto process
labels.append(rand_file) #Append the correct answer tocompare later
return cases, labels
def main():
#Generate cases
data, target = generate_cases(10)
#Run inference on our generated cases doingpreprocessing and postprocessing
internally
results = mnist_engine.infer(data)[0] #Data isreturned in a list by output
layer
#Validate results
correct = 0
print ("[LABEL] | [RESULT]")
for l in range(len(target)):
print (" {} | {} ".format(target[l],results[l]))
if target[l] == results[l]:
correct += 1
print ("Inference: {:.2f}%Correct".format((correct / len(target)) * 100))
使用Lite engine里的engine.save(path)保存TensorRTengine。在使用时将engine文件的位置传到PLAN文件中,如果必要的话,还有前处理与后处理函数。