整体加速过程
1.将tensorflow训练生成的.h5模型结构权重文件转换成.pb的模型文件。
2.将.pb模型文件转换成uff格式的文件并进行解析,同时生成TensorRT的engine。
3.调用生成的engine文件,实现推理加速。
准备工作
1.生成.pb的模型文件
首先我们需要从保存模型的chekpoint文件中,生成.pb的模型文件。
import tensorflow.compat.v1 as tf1
tf1.reset_default_graph()
tf1.keras.backend.set_learning_phase(0) # 调用模型前一定要执行该命令
tf1.disable_v2_behavior() # 禁止tensorflow2.0的行为
# 加载hdf5模型,一定要是将模型结构和权重同时保存的文件。
hdf5_pb_model = tf1.keras.models.load_model('D:/Graduate student/Deep learning/project/crowd dense/MSCNN-master/models/bet_model_weights.h5')
def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):
# tf.graph是创建一张新图,这是指获取图
graph = session.graph
# 上下文管理器 能够覆盖掉默认的图
with graph.as_default():
output_names = output_names or []
print("output_names", output_names)
# 返回此图的序列化图形化表示,node结点,node节点的属性包括name,op,input
# GraphDef保存了从输入层到输出层的计算过程
input_graph_def = graph.as_graph_def()
print(type(input_graph_def))
print("len node1", len(input_graph_def.node))
if clear_devices:
for node in input_graph_def.node:
node.device = ""
# 计算图中的变量取值以常量的形式保存,output_name用来指定保存的节点
frozen_graph = tf1.graph_util.convert_variables_to_constants(session, input_graph_def, output_names)
# 删除不需要进行推理(inference)的节点,在训练过程中,有一些节点,比如Identity和checknerics,它们只在训练时有用,并且可以在仅用于推理的图形中删除
outgraph = tf1.graph_util.remove_training_nodes(frozen_graph) # 云掉与推理无关的内容
for node in outgraph.node:
print('node:', node.name)
print("len node1", len(outgraph.node))
return outgraph
output_folder2 = 'D:/Graduate student/Deep learning/project/crowd dense/MSCNN-master/models'
# 模型冻结,返回TensorFlow会话,如果在创建Session时没有指定Graph,则该Session会加载默认Graph
frozen_graph = freeze_session(tf1.compat.v1.keras.backend.get_session(),
output_names=[out.op.name for out in hdf5_pb_model.outputs])
# 转换成pb模型
tf1.train.write_graph(frozen_graph, output_folder2, "classify.pb", as_text=False) #保存图结构
- 导入必要的库
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
1.pb转uff 并解析模型
# 建立日志模块
G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.INFO)
uff_model = uff.from_tensorflow_frozen_model(FROZEN_GDEF_PATH, OUTPUT_NODE)
parser = uffparser.create_uff_parser()
parser.register_input(INPUT_NODE, INPUT_SIZE, 0)
parser.register_output(OUTPUT_NODE)
这里做的事情是将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时,请一定确认图像的格式。