导出支持多输入尺寸的TF模型至TFLite的标准方法确认

我需要将TensorFlow模型导出为TFLite格式，目标是在移动设备GPU上支持不同输入形状的推理。

方法1：基于固定初始形状导出，后续动态调整 #

我最初使用以下方法导出模型：

# tf model class
class MyModel(tf.keras.models.Model):
...


# Util functions
def build_graph(model, input_shape):
    x = tf.keras.layers.Input(shape=input_shape)
    model =  tf.keras.models.Model(inputs=x, outputs=model(x))
    model.compile(optimizer='adam', loss='binary_crossentropy')
    return model

def save_tflite_model(output_model_path, tflite_model):
    with open(output_model_path, 'wb') as f:
        f.write(tflite_model)
    
def convert_model_from_concrete(model_path, output_model_path, input_shape=(1, None, None, 3)):
    model = tf.saved_model.load(model_path)
    concrete_func = model.signatures[
        tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
    
    concrete_func.inputs[0].set_shape(input_shape)
    
    converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
    converter.experimental_new_converter = True
    
    converter.target_spec.supported_ops = [
        tf.lite.OpsSet.TFLITE_BUILTINS,
        tf.lite.OpsSet.SELECT_TF_OPS 
        ]
    tflite_model = converter.convert()
    print(tf.lite.experimental.Analyzer.analyze(model_content=tflite_model, gpu_compatibility=True))
    
    save_tflite_model(output_model_path, tflite_model)


#Code for exporting my model to TFLite
model = MyModel()
temp_input_shape = (256, 256, 3)
model = build_graph(model, temp_input_shape)
model.save("my_model_256")
convert_model_from_concrete("my_model_256","my_model_256.tflite")

该方法导出的模型可本地加载并运行任意有效形状：

interpreter = tf.lite.Interpreter("my_model_256.tflite")
custom_shape = [1, 512, 512, 3]
interpreter.resize_tensor_input(interpreter.get_input_details()[0]['index'], custom_shape)
interpreter.allocate_tensors()
input =  numpy.random.rand(*custom_shape).astype(np.float32)
input_details = interpreter.get_input_details()
interpreter.set_tensor(input_details[0]['index'], input)
interpreter.invoke()

并且该模型在TFLite Benchmark工具的GPU模式下无报错，可正常运行不同形状的输入。

方法2：直接导出动态输入形状模型 #

参考相关建议，我尝试直接导出带有动态输入尺寸的模型：

#Code for exporting my model to TFLite with dynamic input shape
model = MyModel()
temp_input_shape = (None, None, 3)
model = build_graph(model, temp_input_shape)
model.save("my_model_dynamic")
convert_model_from_concrete("my_model_dynamic","my_model_dynamic.tflite")

该模型在本地运行正常，但使用TFLite Benchmark工具配合GPU delegate，并指定自定义输入形状--use_gpu=true --input_layer=input_8 --input_layer_shape=1, 512, 512, 3时，出现如下错误：

INFO: STARTING!
INFO: Log parameter values verbosely: [0]
INFO: Min num runs: [1]
INFO: Num threads: [8]
INFO: Graph: [./my_model_dynamic.tflite]
INFO: Input layers: [input_1]INFO: Input shapes: [1, 512, 512, 3]
INFO: #threads used for CPU inference: [8]
INFO: Use gpu: [1]
INFO: Loaded model ./my_model_dynamic.tflite
INFO: Initialized TensorFlow Lite runtime.
INFO: Created TensorFlow Lite delegate for GPU.
INFO: GPU delegate created.
VERBOSE: Replacing 343 out of 343 node(s) with delegate (TfLiteGpuDelegateV2) node, yielding 1 partitions for the whole graph.
ERROR: Failed to allocate device memory (clCreateSubBuffer): Invalid buffer size
ERROR: Falling back to OpenGL
INFO: Initialized OpenGL-based API.
ERROR: TfLiteGpuDelegate Init: Shapes are not equal
INFO: Created 0 GPU delegate kernels.
ERROR: TfLiteGpuDelegate Prepare: delegate is not initialized
ERROR: Node number 343 (TfLiteGpuDelegateV2) failed to prepare.
ERROR: Restored original execution plan after delegate application failure.
ERROR: Failed to apply GPU delegate

疑问 #

1. 上述两种导出方法中，哪种是TensorFlow官方支持的？
2. 我修改了Benchmark工具代码使第二种方法可用，但不确定这样的修改是否合法？

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