TensorRT是nvidia家的一款高性能深度学习推理SDK。此SDK包含深度学习推理优化器和运行环境,可为深度学习推理应用提供低延迟和高吞吐量。在推理过程中,基于TensorRT的应用程序比仅仅使用CPU作为平台的应用程序要快40倍。
TensorRT包含两个阶段:编译build和部署deploy。
为了优化你的推理模型,TensorRT将接受你的网络定义,执行优化,包括特定平台优化,并且生成一个推理引擎(inference engine)。这个过程被视作编译阶段(build phase)。编译计算可能耗费相当多的时间,尤其是在嵌入式平台中运行时。因此,一个典型的应用将会构建一个引擎,然后将其序列化为一个plan 文件,以供后续使用。(生成的plan文件并不能够跨平台/TensorRT版本移植)
编译阶段在图层中执行如下优化:
C++ API和Python API在支持开发者的需求方面非常接近。在任何性能是关键的场景中,和在安全性非常重要的情况下,C++应该被使用。Python API 的主要优点是可以使用Python的各种库文件对数据预处理和后处理。
将一个训练好的模型部署到TensorRT上的流程为:
# 导入TensorRT
import tensorrt as trt
# 日志接口,TensorRT通过该接口报告错误、警告和信息性消息
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
使用解析器(parser)导入一个模型,需要执行以下步骤:
import tensorrt as trt
# 以CaffeParse为例
datatype = trt.float32 # 定义数据类型
# 定义配置文件和参数模型路径
deploy_file = 'data/mnist/mnist.prototxt'
model_file = 'data/mnist/mnist.caffemodel'
# 创建builder, network 和 parser
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.CaffeParser() as parser:
model_tensors = parser.parse(deploy=deploy_file, model=model_file, network=network, dtype=datatype)
注: builder必须在network之前创建。不同的解析器有不同的机制标记网络输出。
builde的功能之一是通过搜索CUDA内核目录以获得可用的最快实现,因此有必要使用相同的GPU进行构建,就像优化引擎将在其上运行一样。
IBuilderConfig有很多属性,你可以设置这些属性来控制网络运行的精度,以及自动调整参数等等。其中一个特别重要的属性是 maximum workspace size。
# 使用build对象建造engine
with trt.Builder(TRT_LOGGER) as builder, builder.create_builder_config() as config:
# 当构建一个优化引擎时,这决定了builder可用的内存量,通常应当设置为尽可能高
config.max_workspace_size = 1 << 20
with builder.build_engine(network, config) as engine:
序列化,意味着将engine转化为一种可以存储的格式并且在以后可以进行推理。用于推理使用时,只需要简单地反序列化engine。序列化和反序列化都是可选的。由于从网络定义中创建一个engine是非常耗时的,通常序列化一次并且在推理时反序列化即可。因此,在构建engine之后,用户通常希望序列化它以供以后使用。
# 序列化模型到modelstream
serialized_engine = engine.serialize()
# 反序列化modelstream用于推理。反序列化需要创建runtime对象。
with trt.Runtime(TRT_LOGGER) as runtime:
engine = runtime.deserialize_cuda_engine(serialized_engine)
# 序列化engine并且写入一个file中
with open(“sample.engine”, “wb”) as f:
f.write(engine.serialize())
# 从文件中读取engine并且反序列化
with open(“sample.engine”, “rb”) as f, trt.Runtime(TRT_LOGGER) as runtime:
engine = runtime.deserialize_cuda_engine(f.read())
# engine有一个输入binding_index=0和一个输出binding_index=1
h_input = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(0)), dtype=np.float32)
h_output = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(1)), dtype=np.float32)
# 为输入和输出分配内存
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
# 创建一个流在其中复制输入/输出并且运行推理
stream = cuda.Stream()
#创建一些空间来存储中间激活值。由于引擎包含网络定义和训练参数,因此需要额外的空间。它们被保存在执行上下文中。
with engine.create_execution_context() as context:
# 将输入数据转换到GPU上
cuda.memcpy_htod_async(d_input, h_input, stream)
# 运行推理
context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
# 从GPU上传输预测值
cuda.memcpy_dtoh_async(h_output, d_output, stream)
# 同步流
stream.synchronize()
# 返回主机输出
return h_output
import pycuda.driver as cuda
import pycuda.autoinit
import tensorrt as trt
import common
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
# 分配主机和设备缓冲区,创建流
def allocate_buffers(engine):
h_input = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(0)), dtype=np.float32)
h_output = cuda.pagelocked_empty(trt.volume(context.get_binding_shape(1)), dtype=np.float32)
d_input = cuda.mem_alloc(h_input.nbytes)
d_output = cuda.mem_alloc(h_output.nbytes)
stream = cuda.Stream()
return h_input, h_output, d_input, d_output, stream
def do_inference(context, h_input, h_output, d_input, d_output, stream):
cuda.memcpy_htod_async(d_input, h_input, stream)
context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
cuda.memcpy_dtoh_async(h_output, d_output, stream)
stream.synchronize()
def build_engine_caffe(model_file, deploy_file):
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.CaffeParser() as parser:
builder.max_workspace_size = common.GiB(1)
model_tensors = parser.parse(deploy=deploy_file, model=model_file, network=network, dtype=trt.float32)
network.mark_output(model_tensors.find(ModelData.OUPUT_NAME))
return builder.build_cuda_engine(network)
def Load_normalized_test_case(test_image, pakelocked_buffer):
# 将输入图像转化为一个CHW numpy数组
def main():
caffe_model_file, caffe_deploy_file
with build_engine_caffe(caffe_model_file, caffe_deploy_file) as engine:
h_input, h_output, d_input, d_output, stream = allocate_buffers(engine)
with engine.create_execution_context() as context:
test_image
Load_normalized_test_case(test_image, h_input)
do_inference(context, h_input, h_output, d_input, d_output, stream)
import tensorrt as trt
# 使用python设置层精度 Layer Precision
# 用precision指定层精度
layer.precision = trt.int8
# 设置输出张量数据类型与层实现一致
layer.set_output_type(out_tensor_index, trt.int8)
# builder强制准寻设置的精度
builder.strict_type_constraints = true
# 使用Python使得推理以FP16精度运行
builder.fp16_mode = True
# 通过设置builder标志强制16-bit精度
builder.strict_type_constraints = True
# 通过设置builder标志运行INT8精度模式
builder.int8_mode = True
# 使用Python设置每一层的动态范围
# 为了能够以INT8精度执行推理,必须为每一网络张量设置动态范围。可以使用各种方法导出动态范围值,包括量化感知训练或者仅仅简单地记录上一个训练epoch期间每个张量的最小和最大值
layer = network[layer_index]
tensor = layer.get_output(output_index)
tensor.dynamic_range = (min_float, max_float)
# 同样需要为网络输入设置动态范围
input_tensor = network.get_input(input_index)
input_tensor.dynamic_range = (min_float, max_float)
# INT8校准
# 和测试/验证文件相似,使用输入文件集合作为一个矫正文件数据集。确保校准文件能够代表整个推断数据文件。为了使得TensorRT能够使用校准文件,需要创建batchstream对象。一个bacthstream对象通常被用来配置校准器。
NUM_IMAGES_PER_BATCH = 5
batchstream = ImageBatchStream(NUM_IMAGES_PER_BATCH, calibration_files)
# 用输入节点名称和batch stream创建一个Int8_calibrator对象
Int8_calibrator = EntropyCalibrator(["input_node_name"], batchstream)
# 设置INT8模式和INT校准
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = Int8_calibrator
# 使用Python API创建一个明确的精度网络,需要将EXPLICIT_PRECISION标志传送给builder
network_creation_flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_PRECISION)
self.network = self.builder.create_network(network_creation_flag)
TensorRT 允许开发者可以import、calibrate、generate和deploy优化网络。网络可以直接从Caffe框架中imported,或者通过UFF/ONNX格式从其他框架中imported;网络也可以通过直接设置参数和权重实例化各层以编程地方式创建。
TensorRT为所有支持平台提供了C++实现,以及在x86、aarch64和ppc64le平台上提供Python支持。
终于恰饭了,感谢金主深蓝学院
,如果您想系统性学习TensorRT,可访问深度神经网络加速cudnn和TensorRT课程。
[1]. What’s the Difference Between Deep Learning Training and Inference?
[2]. TensorRT(1)-介绍-使用-安装
[3]. TensorRT Homepages
[4]. TensorRT Developer Guide