[TVM] Падает компиляция модели TensorFlow

Версии пакетов:

tensorflow==2.14.0
apache-tvm==0.14.dev214
tflite==2.1.0

Модель: densenet-121-tf из open_model_zoo в формате saved_model Проблема: Для компиляции модели TF в представление TVM нужно передать граф вычислений. Его получаю 2 способами:

Функциями из репозитория, которые находятся в dl-benchmark/src/model_converters/tensorflow_common.py, получается следующий код:

def _get_tf_model(self, model_path):
    input_op_name = get_input_operation_name(self.args['input_name'])
    output_names = self.args['output_names']
    model, _ = load_model(Path(model_path),
                          input_names=input_op_name,
                          output_names=output_names,
                          const_inputs=[],
                          log=log)
    signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
    func = model._backref_to_saved_model.signatures[signature_key]
    graph = func.graph
    return graph

def _convert_model_from_framework(self, target, dev):
    model_path = self.args['model_path']
    opt_lev = self.args['opt_level']
    model_tf = self._get_tf_model(model_path)
    shape_dict = {self.args['input_name']: self.args['input_shape']}
    log.info('Creating graph module from TensorFlow model')
    model, params = tvm.relay.frontend.from_tensorflow(model_tf, layout=None, shape=shape_dict)
    with tvm.transform.PassContext(opt_level=opt_lev):
        lib = tvm.relay.build(model, target=target, params=params)
    module = tvm.contrib.graph_executor.GraphModule(lib['default'](dev))
    return module

Где функция load_model лежит в dl-benchmark/src/model_converters/tensorflow_common.py и по смыслу возвращает модель или из .pb, или из saved_model(директория). Получаю следующую ошибку:

File "/home/vanya/projects/dl-benchmark/src/inference/./inference_tvm_tensorflow.py", line 195, in main
graph_module = converter.get_graph_module()
File "/home/vanya/projects/dl-benchmark/src/inference/tvm_auxiliary.py", line 33, in get_graph_module
self.graph = self._convert_model_from_framework(target, dev)
File "/home/vanya/projects/dl-benchmark/src/inference/./inference_tvm_tensorflow.py", line 170, in _convert_model_from_framework
model, params = tvm.relay.frontend.from_tensorflow(model_tf, layout=None, shape=shape_dict)
File "/home/vanya/anaconda3/envs/dl_benchmark_env/lib/python3.9/site-packages/tvm/relay/frontend/tensorflow.py", line 1263, in from_tensorflow
mod, params = g.from_tensorflow(graph, layout, shape, outputs)
File "/home/vanya/anaconda3/envs/dl_benchmark_env/lib/python3.9/site-packages/tvm/relay/frontend/tensorflow.py", line 663, in from_tensorflow
func = self._get_relay_func(graph, layout=layout, shape=shape, outputs=outputs)
File "/home/vanya/anaconda3/envs/dl_benchmark_env/lib/python3.9/site-packages/tvm/relay/frontend/tensorflow.py", line 469, in _get_relay_func
missing_operators = self._parse_import_prerequisites(graph)
File "/home/vanya/anaconda3/envs/dl_benchmark_env/lib/python3.9/site-packages/tvm/relay/frontend/tensorflow.py", line 676, in _parse_import_prerequisites
for node in graph.node:
AttributeError: 'FuncGraph' object has no attribute 'node'

По документации TF не понятно, как из FuncGraph получить объект, который будет иметь атрибут node.

Способ из примера TVM (https://tvm.apache.org/docs/how_to/compile_models/from_tensorflow.html#sphx-glr-how-to-compile-models-from-tensorflow-py). Код:

with tf_compat_v1.gfile.GFile([model_path](https://docs.python.org/3/library/stdtypes.html#str), "rb") as f:
graph_def = tf_compat_v1.GraphDef()
graph_def.ParseFromString(f.read())
graph = tf.import_graph_def(graph_def, name="")
# Call the utility to import the graph definition into default graph.
graph_def = tf_testing.ProcessGraphDefParam(graph_def)
# Add shapes to the graph.
with tf_compat_v1.Session() as sess:
    graph_def = tf_testing.AddShapesToGraphDef(sess, "softmax")

Здесь компиляция просто уходит в вечный цикл, причем последние 2 строчки не всегда могут обработаны, потому что не у всех моделей есть сущность softmax

itlab-vision / dl-benchmark

[TVM] Падает компиляция модели TensorFlow #428