tensorflow-onnx icon indicating copy to clipboard operation
tensorflow-onnx copied to clipboard

Very slow (over 8 minutes) when converting keras UpSampling1D layer to ONNX given a tensor with large tensor shape

Open maybeLee opened this issue 3 years ago • 0 comments

Describe the bug Not sure if it is expected behavior, I was converting a densenet-like model trained by Keras and I notice that it takes tf2onnx over 18 minutes to convert this model to onnx. After some attempts, I simplify this performance issue to a keras model with only one layer: UnSampling1D. The triggering condition is to give this UpSampling1D a tensor with a large shape: (batch_size, 20000, 1)

The script to build the model is as follows:

import keras
new_input = keras.layers.Input((20000, 1))
layer_stack = [
        keras.layers.UpSampling1D(size=2)
        ]
layer_input = new_input
for layer in layer_stack:
    y = layer(layer_input)
    layer_input = y
new_model = keras.models.Model(new_input, y)
new_model.summary()
new_model.save("new_model.h5")

The script to convert this model to onnx as as follows:

def transform_onnx(model_path: str):
    model_name = model_path.split(".h5")[0]
    onnx_path = f"{model_name}.onnx"
    import keras
    import tensorflow as tf
    tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
    from keras import backend as K
    K.set_learning_phase(0)
    import tf2onnx
    model = keras.models.load_model(model_path)
    print("Finish loading the model, start converting")
    if tf.__version__.split(".")[0] == 1:
        input_shape = model.layers[0].input_shape
    else:
        input_shape = model.layers[0].input_shape[0]
    spec = (tf.TensorSpec(input_shape, tf.float32, name="input"),)
    _, _ = tf2onnx.convert.from_keras(model, input_signature=spec, \
        opset=15, output_path=onnx_path)
    del model
    del _

import datetime
convert_st = datetime.datetime.now()
transform_onnx("new_model.h5")
convert_et = datetime.datetime.now()
print(f"Total converting time is: {convert_et - convert_st} (seconds)")

Result shows that it takes tf2onnx over 8 minutes to convert this single layer on colab environment with GPU

Total converting time is: 0:08:03.582000 (seconds)

To Reproduce Please run above code or follow this link: https://colab.research.google.com/drive/19CiLFWjkVHer0eoJYA3UPAVMj13aX_Yq?usp=sharing

maybeLee avatar Mar 31 '22 09:03 maybeLee