1. TensorRT的模型转换
当我们需要利用NVIDIA的硬件进行加速模型的时候,很多时候会选择将模型转为TensorRT模型,不同的框架下有各种各样的转换方式,如转为uff然后再转TensorRT,或者先转为onnx再转TensorRT,不过一般都比较繁琐,并且需要修改推断代码。
自TensorFlow1.7及更高版本开始,TensorFlow与TensorRT模型混合使用变得更加方便,通常不需要改动原来的TensorFlow模型的推断代码就能够使用TensorRT模型,具体原理可以参考使用TensorRT加速GPU上的TensorFlow推理。
一般来说,需要先把原来的TensorFlow模型转为TensorRT模型,转换的核心代码如下:
trt_graph = trt.create_inference_graph(graph_def, outputs=output_names, max_batch_size=batch_size, max_workspace_size_bytes=workspace_size, precision_mode=precision_mode)
但是往往在这个过程中,会出现譬如以下的报错:
TensorRT node TRTEngineOp_0 added for segment 0 consisting of 139 nodes failed: Internal: Input shapes must be fully defined when in static mode. Please try is_dynamic_op=True (shape was [?,?,?,?]). Fallback to TF...
究其原因,是因为在转换为TensorRT模型的时候,要求输入是确定的shape,如果不是需要将is_dynamic_op参数设置为True,这样将转为动态模式。
2. 静态模式与动态模式区别
2.1. 静态模式
转之前的模型,输入的维度形状都是确定的,不可以变。这种情况下,TensorRT已经将engine离线创建了,在加载TensorRT模型的时候,会自动创建engine。
2.2. 动态模式
转之前的输入维度是动态的(最经常的情况,就是批次的维度不确定,但是其他维度确定),这个时候,在转换的时候其实并没有离线创建engine,但是由于TensorRT在运行的时候又要求所有的维度都是确定的,因此,TensorRT每遇到一个新的形状的时候,就会创建一个新的engine。在实际的推理中,有时候速度会突然变慢,往往是创建engine造成的。
3. engine的重用机制
上面说到,在动态模式下,遇到新的输入形状,TensorRT会新创建一个新engine。但并不是所有情况都会重新创建,TensorRT有其engine重用机制,只要符合下面两个条件即可:
- 当前engine的batch size大于等于新的input的batch size
- engine的其他维度与新的input的其他维度匹配
举个例子。我们模型的输入维度为[?,128,128,3],其中第一个维度为batch size维度。当遇到第一个新input,它的输入形状为[10,128,128,3],这个时候,TensorRT会创建一个新的engine(就叫A吧),它输入的形状为[10,128,128,3],并将该engine缓存。
如果遇到一个新的输入,如[3,128,128,3],这个时候会先去缓存寻找是否有符合的engine,因为batch size比engine A的要小,而且其他维度是一样的,所以就不再创建新的engine,直接利用A进行推理了。
当然,engine的缓存个数是有限的,需要在转换的时候进行设置(maximum_cached_engines参数 ),一般建议是5.
最后修改完的参数为:
trt_graph = trt.create_inference_graph(graph_def, outputs=output_names, max_batch_size=batch_size, is_dynamic_op=True, max_workspace_size_bytes=workspace_size, maximum_cached_engines=5, precision_mode=precision_mode)
4. 参考
(完)