TensorRT中converter.build()的作用是什么?

3

TensorRT 官方文档 列出了两种将 TensorFlow SavedModel 转换为 TensorRT SavedModel 的方法:第一种是

from tensorflow.python.compiler.tensorrt import trt_convert as trt
converter = trt.TrtGraphConverterV2(input_saved_model_dir=input_saved_model_dir)
converter.convert()
converter.save(output_saved_model_dir)

and the second is

import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt

conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS
conversion_params = conversion_params._replace(
    max_workspace_size_bytes=(1<<32))
conversion_params = conversion_params._replace(precision_mode="FP16")
conversion_params = conversion_params._replace(
    maximum_cached_engiens=100)

converter = trt.TrtGraphConverterV2(
    input_saved_model_dir=input_saved_model_dir,
    conversion_params=conversion_params)
converter.convert()
def my_input_fn():
  for _ in range(num_runs):
    Inp1 = np.random.normal(size=(8, 16, 16, 3)).astype(np.float32)
    inp2 = np.random.normal(size=(8, 16, 16, 3)).astype(np.float32)
    yield inp1, inp2
converter.build(input_fn=my_input_fn)
converter.save(output_saved_model_dir)

saved_model_loaded = tf.saved_model.load(
    output_saved_model_dir, tags=[tag_constants.SERVING])
graph_func = saved_model_loaded.signatures[
    signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
frozen_func = convert_to_constants.convert_variables_to_constants_v2(
    graph_func)
output = frozen_func(input_data)[0].numpy()

除去所有的导入、推断等样板代码后,似乎差异在于对converter.build()的调用。文档对此函数的解释如下:

"该方法通过构建TensorRT引擎来优化转换后的函数(由convert()返回)。这在用户想要在运行之前执行优化时非常有用。通过使用从参数input_fn接收到的输入数据运行convert()函数进行优化。该参数是一个生成器函数,以列表或元组的形式生成输入数据。"

在这个上下文中,“在运行之前”是什么意思?这些“优化”是在模型加载时、在第一次推理时还是在每次使用转换后的模型进行推理时执行的?那些“优化”究竟是什么?将模型转换为TensorRT本身不就已经是优化了吗?

我之所以问这个问题,是因为如果我调用converter.build(),转换似乎会以不可预测的方式失败,并且需要很长时间(超过两个小时)才能运行,而且没有产生任何有意义的输出,所以我想知道如果不调用它,我会失去多少,以及是否有更全面的文档介绍如何在TensorRT中使用TF2.x SavedModels。

感谢事先回答的人!!


哇,nVidia的文档中有一个明显的拼写错误:Inp1inp1。他们不可能没有注意到这个代码片段的问题。 - couka
1个回答

6
根据我的理解(阅读 TensorFlow文档后),converter.convert()函数将图形转换为tf-trt,用TRTEngineOp替换它所能替换的节点,但不会创建用于推断的实际引擎文件。调用converter.build()函数会创建引擎文件,但是这些文件是基于提供给input_fn的输入大小和构建所在平台而生成的。因此,不调用converter.build()的原因可能是您事先不知道输入形状,或者无法在要运行推断的平台上进行构建。我很难想象每次进行前向传递时都会创建新的引擎文件,但肯定会为每个新的输入形状创建一个引擎文件。从示例中不清楚input_fn中的输入是否以其他方式使用,但如果返回不同形状的输入,则应为每个输入大小创建一个引擎文件。
作为示例,提供以下输入函数将生成输入大小为(112, 112, 3)和(224, 224, 3)的两个引擎:
def input_fn():
  input_sizes = [[112, 112], [224, 224]]
  for size in input_sizes:
    inp1 = np.random.normal(size=(1, *size, 3)).astype(np.float32)
    yield [inp1]

关于您的input_fn,您是否将两个图像作为网络的输入?对我而言,像上面示例中那样以列表形式返回单个图像可行(元组由于某些原因无法使用,尽管文档中说应该可以)。
希望这能有所帮助。

1
谢谢,我没有意识到我必须在我将运行推理的平台上运行它。对我来说,这是Jetson Nano。一开始每次都会耗尽GPU内存。将max_workspace_size_bytes=(1<<32)减小到max_workspace_size_bytes=(1<<25)对我很有帮助。 - couka

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接