如何将 TensorRT SavedModel 加载到 TensorFlow Estimator？

Question

I'm using TensorFlow 1.14, and to load a TensorFlow SavedModel to an Estimator the following code works for me:

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

However, when I used TensorRT to convert the TensorFlow SavedModel to TensorRT SavedModel, it returns an error message:

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

I have looked at it further, and it looks like the problem is that TensorRT does not generate any variables information (including variables.index) in the SavedModel directory, which makes the above error. Does anyone know how to resolve this problem?

Answer 1

For anyone who is interested, below is the solution I have figured out by myself: Normally, a TF SavedModel can be loaded to Estimator using:

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

However, when loading TensorRT SavedModel errors occur because TensorRT converts all variables to constants, thus there is no variables' information in the SavedModel directory (eg no variables.index) → errors occur since Estimator try to load variables. Steps to fix the problem:

• We need to go to file: "/usr/local/lib/python3.6/dist-packages/tensorflow_estimator/python/estimator/estimator.py", line 2330, in _get_saved_model_ckpt and comment out the check for 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)"""

• Go to file: "/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)] and make changes so that Estimator does not try to load variables from 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))

• Go to file: "/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) and comment out the check for "global_step" in case the model is generated from NVIDIA examples (thus no training was done and "global_step" was not set):

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

• Go to file: "/usr/local/lib/python3.6/dist-packages/tensorflow/python/training/checkpoint_utils.py", line 291, in init_from_checkpoint init_from_checkpoint_fn) and just put return at the beginning of the init_from_checkpoint function so that it does not try to load the checkpoint:

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

The loading process should be fine after all the above changes.