如何将 TensorRT SavedModel 加载到 TensorFlow Estimator？

Question

我正在使用 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），导致了上面的错误。 有谁知道如何解决这个问题？

Answer 1

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

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

但是，在加载 TensorRT SavedModel 时会发生错误，因为 TensorRT 将所有变量都转换为常量，因此 SavedModel 目录中没有变量信息（例如，没有 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

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