tensorflow打印模型图-尊龙游戏旗舰厅官网
使用tensorrt加速tensorflow模型的推理应该是很有市场的一种应用了,但是使用python的、易懂的例子并不多,官方的文档在这方面也是很不友好。
所以,本文旨在提供一个能把原理讲明白,代码能跑的起来的实例,本例中用到模型是inception v3
准备工作
首先我们需要从保存模型的chekpoint文件中,生成.pb的模型文件。这一步叫做模型的持久化,具体的做法可以参考之前写的这篇文章:
春天不是读书天:[深度学习] tensorflow中模型的freeze_graph
2. 导入必要的库
import tensorflow as tf import uff import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit from tensorrt.parsers import uffparser- uff:是将刚才的pb转化为tensorrt引擎支持的uff文件,该文件可以序列化,也可以直接当作流传过去。
- pycyda:用于显卡cuda编程的,如果要使用tensorrt的python api,这是一个必须的库
- uffparser :用于解析uff模型
3. 参数设置
model_dir = './model_seg/model.pb' channel = 3 height = 299 width = 299 engine_path = './model_seg/model_.pb.plan' input_node = 'input' output_node = ['inceptionv3/logits/spatialsqueeze'] input_size = [channel, height ,width] max_batch_size = 1 max_workspace = 1<<30- model_dir:第一步中生成的pb模型地址
- channel、height、width:图片的通道、高和宽,根据模型的输入大小确定
- engine_path:等会保存tensorrt引擎的地址
- input_node:模型的输入节点
- output_node:模型的输出节点,是一个列表,如果有许多个输出节点,就将节点名都列入这个列表中
- input_size:输入图片的大小,注意通道在前还是后,这里输入的是 channel, height ,width
- max_batch_size:在推理的时候,每次输入几张图片
- max_workspace:显存的大小1<<30也就是1gb的大小。有的时候,程序运行是会报内存溢出的错,这个时候就可以调小max_workspace,比如2 << 10
将tensorflow模型转换成tensorrt
这里做的事情是将pb的文件格式转成了uff文件格式。你需要知道的一个概念是,uff(universal framework format)是一种描述dnn执行图的数据格式。绑定执行图的是输入与输出,所以parser.register_input和parser.register_output做的事情是将tensorflow模型的输入输出在uff文件中记录。
注意,对于多个输出,因为output_node是一个列表,所以将多个输出节点依次放入列表就可以了。
如果是多个输入的话,则需要将输入节点名一个个的记录在uff中。register_input()需要3个参数:
- name – input name.
- shape – input shape.
- order – input order on which the framework input was originally.
假设你的模型在输入层同时输入了三张图片,那么你需要定义3个输入节点,并且指定order分别为0、1、2。这里的order指的是模型的输入在uff结构中的顺序,这种order在接下来的binding会得到体现。
parser.register_input(input_node1, input_size, 0) parser.register_input(input_node2, input_size, 1) parser.register_input(input_node3, input_size, 2)2. 保存模型
engine = trt.utils.uff_to_trt_engine(g_logger,uff_model,parser,max_batch_size,max_workspace,datatype=trt.infer.datatype.float)以上代码创建了tensorrt中的engine,即引擎,这个engine将负责模型的前向运算。tensorrt是一个用于推理的加速工具,所以前向计算就够了。
在engine创建成功之后,就可以使用了。不过,一个建议是将结果保存下来。毕竟到目前为止,虽然代码很少,但是将pb文件成功转换成uff文件是不容易的(谁用谁知道!)
使用以下语句,我们就保存了一个.plan文件。plan文件是运行引擎用于执行网络的序列化数据。包含权重,网络中执行步骤以及用来决定如何绑定输入与输出缓存的网络信息。
trt.utils.cwrite_engine_to_file('./model_.pb.plan',engine.serialize())使用tensorrt实现推理
现在,让我们调用之前保存的plan文件,启用引擎,开始使用tensorrt实现推理。
engine = trt.utils.load_engine(g_logger, './model_.pb.plan')引擎叫做engine,而引擎运行的上下文叫做context。engine和context在推理过程中都是必须的,这两者的关系如下:
context = engine.create_execution_context() engine = context.get_engine()在运行前向运算前,我们还需要做一次确认。get_nb_bindings()是为了获取与这个engine相关的输入输出tensor的数量。对于本例中单输入输出的模型,tensor的数量是2。如果有多个输入输出,这个确认值就要相应的变化,比如3个输入,1个输出的模型,tensor的数量就是4。我们需要知道这个数量,是为了之后的显存分配做准备。
print(engine.get_nb_bindings()) assert(engine.get_nb_bindings() == 2)现在准备好一张可以输入给模型的图像 img.jpg,并且转换成fp32
img = cv2.imread(img.jpg) img = img.astype(np.float32)同时,创建一个array来“接住”输出数据。为什么说“接住”呢,因为之后你就会看到,引擎做前向推理计算的时候,是生成了一个数据流,这个数据流会写入output array中
#create output array to receive data output_size = 10 output = np.zeros(output_size , dtype = np.float32)我们需要为输入输出分配显存,并且绑定。
# 使用pycuda申请gpu显存并在引擎中注册 # 申请的大小是整个batchsize大小的输入以及期望的输出指针大小。 d_input = cuda.mem_alloc(1 * img.size * img.dtype.itemsize) d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)# 引擎需要绑定gpu显存的指针。pycuda通过分配成ints实现内存申请。 bindings = [int(d_input), int(d_output)]现在,我们可以开始tensorrt上的推理计算了!
# 建立数据流 stream = cuda.stream() # 将输入传给cuda cuda.memcpy_htod_async(d_input, img, stream) # 执行前向推理计算 context.enqueue(1, bindings, stream.handle, none) # 将预测结果传回 cuda.memcpy_dtoh_async(output, d_output, stream) # 同步 stream.synchronize()这个时候,如果你将output打印出来,就会发现output数组中已经有值了,这就是tensorrt计算的结果。
如过你使用tensorflow的方法,对同一组输入数据做预测,看看计算的结果是否一致 ,因为精度的差异会有一些差异,但是大体上来说,使用tensorflow和tensorrt,会得到一致的结果。
特别注意!
tensorrt和tensorflow的数据格式不一样,tensorflow是nhwc格式,即channel_last,而tensorrt中是nchw格式,即channel_first,比如一张rgb图像,在tensorflow中表示为(224, 224, 3),在tensorrt中就是(3,224, 224)。所以使用tensorrt时,请一定确认图像的格式。
参考资料:
tensorrt加速tensorflow模型推理(inception v3为例)
https://blog.csdn.net/abrams90/article/details/80410308
总结
以上是尊龙游戏旗舰厅官网为你收集整理的tensorflow打印模型图_[深度学习]tensorrt加速tensorflow实例的全部内容,希望文章能够帮你解决所遇到的问题。
- 上一篇:
- 下一篇: