如何将 TensorRT SavedModel 加载到 TensorFlow Estimator？

Posted 2023-02-16

【中文标题】如何将 TensorRT SavedModel 加载到 TensorFlow Estimator？

【英文标题】：How to load a TensorRT SavedModel to a TensorFlow Estimator?

【发布时间】：2020-09-17 09:20:10

【问题描述】：

我使用的是 TensorFlow 1.14，要将 TensorFlow SavedModel 加载到 Estimator，以下代码适用于我：

estimator = tf.contrib.estimator.SavedModelEstimator(saved_model_dir)
prediction_results = estimator.predict(input_fn)

但是，当我使用 TensorRT 将 TensorFlow SavedModel 转换为 TensorRT SavedModel 时，它会返回错误消息：

ValueError: Directory provided has an invalid SavedModel format: saved_models/nvidia_fp16_converted

我仔细看了下，貌似问题是TensorRT在SavedModel目录下没有生成任何变量信息（包括variables.index），导致出现上述错误。有谁知道如何解决这个问题？

【问题讨论】：

为什么要使用 TF 估计器来推断 TensorRT 引擎？

这是我必须开发的解决方案的要求，因为解决方案的其他部分已经在使用 TF 估计器。

【参考方案1】：

对于任何有兴趣的人，以下是我自己想出的解决方案： 通常，可以使用以下方法将 TF SavedModel 加载到 Estimator：

estimator = tf.contrib.estimator.SavedModelEstimator(SAVED_MODEL_DIR)

但是，在加载 TensorRT SavedModel 时会发生错误，因为 TensorRT 会将所有变量转换为常量，因此 SavedModel 目录中没有变量的信息（例如 no variables.index）→ 由于 Estimator 尝试加载变量而发生错误。解决问题的步骤：

我们需要转到文件："/usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/estimator.py", line 2330, in _get_saved_model_ckpt 并注释掉对 variables.index 的检查

"""if not gfile.Exists(os.path.join(saved_model_utils.get_variables_dir(saved_model_dir),
compat.as_text('variables.index'))):
raise ValueError('Directory provided has an invalid SavedModel format: %s'
% saved_model_dir)"""

转到文件："/usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/canned/saved_model_estimator.py", line 145, in __init__checkpoint_utils.list_variables(checkpoint)] 并进行更改，以使 Estimator 不会尝试从 SavedModel 加载变量：

"""checkpoint = estimator_lib._get_saved_model_ckpt(saved_model_dir) # pylint: disable=protected-access
vars_to_warm_start = [name for name, _ in
checkpoint_utils.list_variables(checkpoint)]
warm_start_settings = estimator_lib.WarmStartSettings(
ckpt_to_initialize_from=checkpoint,
vars_to_warm_start=vars_to_warm_start)"""
warm_start_settings = estimator_lib.WarmStartSettings(ckpt_to_initialize_from = estimator_lib._get_saved_model_ckpt(saved_model_dir))

转到文件："/usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/canned/saved_model_estimator.py", line 256, in _model_fn_from_saved_modeltraining_util.assert_global_step(global_step_tensor) 并注释掉对“global_step”的检查，以防模型是从 NVIDIA 示例生成的（因此未进行任何训练且未设置“global_step”）：

#global_step_tensor = training_util.get_global_step(g)
#training_util.assert_global_step(global_step_tensor)

转到文件："/usr/local/lib/python3.6/dist-packages/tensorflow/python/training/checkpoint_utils.py", line 291, in init_from_checkpoint init_from_checkpoint_fn)，然后将 return 放在 init_from_checkpoint 函数的开头，这样它就不会尝试加载检查点：

def _init_from_checkpoint(ckpt_dir_or_file, assignment_map):
"""See `init_from_checkpoint` for documentation."""
return

完成上述所有更改后，加载过程应该没问题。