TF-TRT转换器转换自有训练模型失败问题求助

环境 #

• 使用镜像：nvcr.io/nvidia/tensorflow:23.12-tf2-py3
• 第三方测试代码可正常运行，但转换自有H5格式训练模型时失败

转换流程代码 #

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

# 加载H5格式训练模型
h5_model_path = 'model/path/h5/model_name'
h5_model = keras.models.load_model(h5_model_path, compile=False)

# 转换为SavedModel格式（TF-TRT要求格式）
saved_model_path = 'model/path/saved_model/model_name'
tf.saved_model.save(h5_model, saved_model_path)

# 初始化TF-TRT转换器
conversion_param = trt.TrtConversionParams(precision_mode=trt.TrtPrecisionMode.FP16)
converter = trt.TrtGraphConverterV2(input_saved_model_dir=saved_model_path, conversion_params=conversion_param)

# 执行转换时触发报错
converter.convert()

报错信息 #

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/training/py_checkpoint_reader.py", line 92, in NewCheckpointReader
    return CheckpointReader(compat.as_bytes(filepattern))
RuntimeError: Unsuccessful TensorSliceReader constructor: Failed to find any matching files for /model/path/saved_model/model_name/variables/variables

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/saved_model/load.py", line 1031, in load_partial
    loader = Loader(object_graph_proto, saved_model_proto, export_dir,
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/saved_model/load.py", line 226, in __init__
    self._restore_checkpoint()
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/saved_model/load.py", line 561, in _restore_checkpoint
    load_status = saver.restore(variables_path, self._checkpoint_options)
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/checkpoint/checkpoint.py", line 1415, in restore
    reader = py_checkpoint_reader.NewCheckpointReader(save_path)
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/training/py_checkpoint_reader.py", line 96, in NewCheckpointReader
    error_translator(e)
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/training/py_checkpoint_reader.py", line 31, in error_translator
    raise errors_impl.NotFoundError(None, None, error_message)
tensorflow.python.framework.errors_impl.NotFoundError: Unsuccessful TensorSliceReader constructor: Failed to find any matching files for /model/path/saved_model/model_name/variables/variables

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/model/code/convert_model.py", line 106, in eval
    converter.convert()
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/compiler/tensorrt/trt_convert.py", line 1453, in convert
    self._saved_model = load.load(self._input_saved_model_dir,
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/saved_model/load.py", line 900, in load
    result = load_partial(export_dir, None, tags, options)["root"]
  File "/usr/local/lib/python3.10/dist-packages/tensorflow/python/saved_model/load.py", line 1034, in load_partial
    raise FileNotFoundError(
FileNotFoundError: Unsuccessful TensorSliceReader constructor: Failed to find any matching files for /model/path/saved_model/model_name/variables/variables
 You may be trying to load on a different device from the computational device. Consider setting the `experimental_io_device` option in `tf.saved_model.LoadOptions` to the io_device such as '/job:localhost'.

已确认事项 #

• SavedModel目录结构与测试代码一致
• /model/path/saved_model/model_name/variables目录下存在variables.data-00000-of-00001和variables.index文件

解决方案尝试 #

1. 指定加载时的IO设备 #

根据报错提示的设备不匹配问题，在初始化TF-TRT转换器时，通过LoadOptions强制指定本地IO设备加载模型：

from tensorflow.python.saved_model import load_options

# 添加加载选项
load_opt = load_options.LoadOptions(experimental_io_device='/job:localhost')
converter = trt.TrtGraphConverterV2(
    input_saved_model_dir=saved_model_path,
    conversion_params=conversion_param,
    input_saved_model_load_options=load_opt
)

2. 验证SavedModel完整性 #

转换前手动加载SavedModel，确认模型本身可正常读取，排除保存环节的问题：

# 加载SavedModel测试
loaded_model = tf.saved_model.load(saved_model_path, options=load_opt)
# 可选：用dummy输入测试推理
dummy_input = tf.random.normal([1, 224, 224, 3])  # 根据模型实际输入维度调整
output = loaded_model(dummy_input)

3. 检查路径权限与绝对路径 #

• 确认容器内对模型路径有完整的读写权限
• 替换相对路径为绝对路径，避免路径解析错误

4. 尝试旧版转换方式 #

若上述方案无效，可尝试兼容旧版的TrtGraphConverter（注意与TF版本匹配）：

converter = trt.TrtGraphConverter(
    input_saved_model_dir=saved_model_path,
    precision_mode=trt.TrtPrecisionMode.FP16
)
converter.convert()

内容的提问来源于stack exchange，提问作者Carpriccio_jh