Проблема с преобразованием модели Tensorflow Superpoint в модель tflite для Android

Я пытаюсь преобразовать сохраненную модель из реализации Superpoint в Tensorflow в модель tflite, чтобы протестировать ее на Android.

Я начал с загрузки сохраненной модели из github: https://github.com/rpautrat/SuperPoint/tree/master/pretrained_models

Модель находится в формате SavedModel. При проверке входов и выходов модели с использованием:

saved_model_cli show --dir sp_v6 --all

Я получаю следующий результат:

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['image'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, -1, -1, 1)
        name: superpoint/image:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['descriptors'] tensor_info:
        dtype: DT_FLOAT
        shape: (1, -1, -1, 256)
        name: superpoint/descriptors:0
    outputs['descriptors_raw'] tensor_info:
        dtype: DT_FLOAT
        shape: (1, -1, -1, 256)
        name: superpoint/descriptors_raw:0
    outputs['logits'] tensor_info:
        dtype: DT_FLOAT
        shape: (1, -1, -1, 65)
        name: superpoint/logits:0
    outputs['pred'] tensor_info:
        dtype: DT_INT32
        shape: (1, -1, -1)
        name: superpoint/pred:0
    outputs['prob'] tensor_info:
        dtype: DT_FLOAT
        shape: (1, -1, -1)
        name: superpoint/prob:0
    outputs['prob_nms'] tensor_info:
        dtype: DT_FLOAT
        shape: (1, -1, -1)
        name: superpoint/prob_nms:0
  Method name is: tensorflow/serving/predict

Насколько мне известно, модели tflite в Android не могут обрабатывать динамические вводы, поэтому я попытался изменить ввод на фиксированный ввод, используя следующий код:

#use tensorflow v1
import tensorflow as tf

def frozen_graph_maker(export_dir,output_graph):
        with tf.Session(graph=tf.Graph()) as sess:
                tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], export_dir)
                output_nodes = ['superpoint/logits', 'superpoint/prob', 'superpoint/descriptors_raw', 'superpoint/descriptors', 'superpoint/prob_nms', 'superpoint/pred']
                output_graph_def = tf.graph_util.convert_variables_to_constants(
                sess, # The session is used to retrieve the weights
                sess.graph_def,
                output_nodes# The output node names are used to select the usefull nodes
                )       
        # Finally we serialize and dump the output graph to the filesystem
        with tf.gfile.GFile(output_graph, "wb") as f:
                f.write(output_graph_def.SerializeToString())

if __name__ == "__main__":
        export_dir='./sp_v6/'
        output_graph = "./frozen_graph.pb"
        frozen_graph_maker(export_dir,output_graph)

В этой игре я замороженный график, в котором я изменяю размер ввода, используя:

#Use tensorflow v2
import tensorflow.compat.v1 as tf
output_graph = "./new_frozen_graph.pb"

def load_frozen_graph(frozen_file='frozen.pb'):
    graph = tf.Graph()
    with graph.as_default():
      od_graph_def = tf.GraphDef()
      with tf.gfile.GFile(frozen_file, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')
    return graph

graph = load_frozen_graph('./frozen_graph.pb')
print('Tensor shapes before import map')
input_tensor = graph.get_tensor_by_name('superpoint/image:0')
print(input_tensor)

new_graph = tf.Graph()
with new_graph.as_default():
    new_input = tf.placeholder(dtype=tf.float32, shape=[1, 320, 320, 1], name='superpoint/image')
    tf.import_graph_def(graph.as_graph_def(), name='', input_map={'superpoint/image:0': new_input}, return_elements=['superpoint/logits:0', 'superpoint/prob:0', 'superpoint/descriptors_raw:0', 'superpoint/descriptors:0', 'superpoint/prob_nms:0', 'superpoint/pred:0'])

with tf.Session(graph=new_graph) as sess:
    output_nodes = ['superpoint/logits', 'superpoint/prob', 'superpoint/descriptors_raw', 'superpoint/descriptors', 'superpoint/prob_nms', 'superpoint/pred']
    output_graph_def = tf.graph_util.convert_variables_to_constants(
    sess, # The session is used to retrieve the weights
    sess.graph_def,
    output_nodes# The output node names are used to select the usefull nodes
    )

with tf.gfile.GFile(output_graph, "wb") as f:
    f.write(output_graph_def.SerializeToString())

Этот замороженный график с измененным размером ввода, затем я конвертирую в формат savedModel, чтобы я мог преобразовать его в формат tflite.

import tensorflow as tf
import os
import shutil
from tensorflow.python import ops

def get_graph_def_from_file(graph_filepath):
  tf.compat.v1.reset_default_graph()
  with ops.Graph().as_default():
    with tf.compat.v1.gfile.GFile(graph_filepath, 'rb') as f:
      graph_def = tf.compat.v1.GraphDef()
      graph_def.ParseFromString(f.read())
      return graph_def

def convert_graph_def_to_saved_model(export_dir, graph_filepath, input_name, outputs):
  graph_def = get_graph_def_from_file(graph_filepath)
  with tf.compat.v1.Session(graph=tf.Graph()) as session:
    tf.import_graph_def(graph_def, name='')
    tf.compat.v1.saved_model.simple_save(
        session,
        export_dir,# change input_image to node.name if you know the name
        inputs={input_name: session.graph.get_tensor_by_name('{}:0'.format(node.name))
            for node in graph_def.node if node.op=='Placeholder'},
        outputs={t.rstrip(":0"):session.graph.get_tensor_by_name(t) for t in outputs}
    )
    print('Graph converted to SavedModel!')

tf.compat.v1.enable_eager_execution()

input_name="superpoint/image"
outputs = ['superpoint/logits:0', 'superpoint/prob:0', 'superpoint/descriptors_raw:0', 'superpoint/descriptors:0', 'superpoint/prob_nms:0', 'superpoint/pred:0']
shutil.rmtree('./saved_model', ignore_errors=True)
convert_graph_def_to_saved_model('./saved_model', './new_frozen_graph.pb', input_name, outputs)

Проверяя входы и выходы этой модифицированной модели, я получаю следующее:

MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['superpoint/image'] tensor_info:
        dtype: DT_FLOAT
        shape: (1, 320, 320, 1)
        name: superpoint/image:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['superpoint/descriptors'] tensor_info:
        dtype: DT_FLOAT
        shape: (1, 320, 320, 256)
        name: superpoint/descriptors:0
    outputs['superpoint/descriptors_raw'] tensor_info:
        dtype: DT_FLOAT
        shape: (1, 40, 40, 256)
        name: superpoint/descriptors_raw:0
    outputs['superpoint/logits'] tensor_info:
        dtype: DT_FLOAT
        shape: (1, 40, 40, 65)
        name: superpoint/logits:0
    outputs['superpoint/pred'] tensor_info:
        dtype: DT_INT32
        shape: unknown_rank
        name: superpoint/pred:0
    outputs['superpoint/prob'] tensor_info:
        dtype: DT_FLOAT
        shape: (1, 320, 320)
        name: superpoint/prob:0
    outputs['superpoint/prob_nms'] tensor_info:
        dtype: DT_FLOAT
        shape: unknown_rank
        name: superpoint/prob_nms:0
  Method name is: tensorflow/serving/predict

Сейчас вход выглядит хорошо, но некоторые выходы имеют неизвестный ранг. Но я пытаюсь продолжить преобразование tflite, как показано ниже:

import tensorflow.lite as lite
saved_model_dir = './saved_model'
converter = lite.TFLiteConverter.from_saved_model(saved_model_dir)
                                    
tflite_model = converter.convert()

with open('./new_frozen.tflite', 'wb') as w:
    w.write(tflite_model)

Но затем во время преобразования я получаю следующую ошибку несоответствия:

tensorflow.lite.python.convert.ConverterError: <unknown>:0: error: type of return operand 3 ('tensor<?x?x?xi32>') doesn't match function result type ('tensor<1x?x?xi32>') in function @main
<unknown>:0: note: see current operation: "std.return"(%32, %30, %36, %47, %40, %45) : (tensor<1x320x320x256xf32>, tensor<1x40x40x256xf32>, tensor<1x40x40x65xf32>, tensor<?x?x?xi32>, tensor<1x320x320xf32>, tensor<?x?x?xf32>) -> ()

Было бы здорово, если бы кто-нибудь мог мне помочь с этим и сказать, что я делаю не так, и есть ли более простой способ преобразовать эту модель тензорного потока в tflite для вывода Android. Я попытался удалить два выхода с неизвестной формой из графика, и затем преобразование сработало, но я хотел бы преобразовать его с этими выходами. Спасибо.