QAT model saving bug: Unable to save function b'__inference_separable_conv2d_layer_call_fn_961'

gcunhase commented 2 years ago

Describe the bug I have a simple model with input layer and a SeparableConv2D layer (note that this issue also happens with Conv2DTranspose). I'm quantizing this model by adding quantize_and_dequantize_v2 nodes at the input and weights of separableconv2d layer (commented in the code). However, I am unable to save the model as SavedModel (whereas converting the Keras model directly to ONNX works):

...
Traceback (most recent call last):
  File "/home/nvidia/PycharmProjects/nvbugs/internal_filed/tf_key_inference_bug/TF_bug_separableconv2d/sample.py", line 24, in <module>
    model.save(model_save_path)
  File "/home/nvidia/PycharmProjects/nvbugs/venv38_trt_regression/lib/python3.8/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/nvidia/PycharmProjects/nvbugs/venv38_trt_regression/lib/python3.8/site-packages/tensorflow/python/saved_model/save.py", line 403, in map_resources
    raise ValueError(
ValueError: Unable to save function b'__inference_separable_conv2d_layer_call_fn_961' because it captures graph tensor Tensor("model/quant_separable_conv2d/LastValueQuant_1/QuantizeAndDequantizeV4:0", shape=(3, 3, 3, 1), dtype=float32) from a parent function which cannot be converted to a constant with `tf.get_static_value`.

System information

TensorFlow version (installed from source or binary): tensorflow-gpu==2.8.0
TensorFlow Model Optimization version (installed from source or binary): N/A
Python version: 3.8.11

Describe the expected behavior Keras model can be saved and loaded as SavedModel.

Describe the current behavior Keras model cannot be saved as SavedModel (loading not able to test since saving is not working).

Code to reproduce the issue Please download the scripts to reproduce from: https://drive.google.com/file/d/1__EimBaQAXIgNPmYKl99uiBLR8FJe2EN/view?usp=sharing

Requirements:

pip install tensorflow-gpu==2.8.0 tf2onnx

Reproduce the issue by running:
```
python sample_qat.py
```

Additional context

This issue also happened with DepthwiseConv2D in tf<2.8.0, but has been fixed in 2.8.0. Please see the previous issue here: https://github.com/tensorflow/model-optimization/issues/868 .
This issue also happens with ConvTranspose2D.

gcunhase commented 2 years ago

@daverim this bug is related to this bug. Please check if possible. Thank you!

gcunhase commented 2 years ago

Any updates?

daverim commented 2 years ago

Hi @gcunhase if you add the quantize_and_dequantize nodes, you will need to create a tf.function that contains both the original function and the graph and save this.

For example if you have a model like my_model(tensor), and you want my_model(quantize_and_dequantize(x)), you will need to encapsulate this in a new model (function) to save, as the original model has not been modified. The easiest way to save is

class NewModel(tf.keras.Model):
  def __init__(self, my_model):
     super(NewModel, self).__init__()
     self.my_model = my_model

  def call(x):
    return self.my_model(tf.quantize_and_dequantize(x))

new_model = NewModel(my_model)
tf.saved_model.save(new_model, '/path/to/saved', )

and then you can save this as a saved model, otherwise the node will not be traceable.

gcunhase commented 2 years ago

@daverim thank you for your reply. I tried your WAR and I'm still getting the same error:

ValueError: Unable to save function b'__inference_sep_conv2d_1_layer_call_fn_1381' because it captures graph tensor Tensor("new_model/toy_separableconv2d_test/quant_sep_conv2d_1/LastValueQuant_1/QuantizeAndDequantizeV4:0", shape=(3, 3, 1, 1), dtype=float32) from a parent function which cannot be converted to a constant with `tf.get_static_value`.

New code snippet:

model = quantize_model(model)
model_save_path = 'weights/sample_qat'

class NewModel(tf.keras.Model):
    def __init__(self, my_model):
        super(NewModel, self).__init__()
        self.my_model = my_model

    def call(self, x):
        return self.my_model(x)

q_model = NewModel(q_model)
q_model(tf.random.normal(shape=(1, 3, 224, 224)))
q_model.save(model_save_path)

Please note that the code snippet I sent when opening this bug works for models with Conv2D, Dense, and many other layers, but fails when I include SeparableConv2D or ConvTranspose2D in the model. It also failed with DethwiseConv2D in tensorflow < 2.8. Can you think of any reason why? Tnank you!

daverim commented 2 years ago

you need to move quantize_model into the call or the init

self.my_model = quantize_model(model)

On Wed, Apr 27, 2022 at 9:45 AM Gwena Cunha @.***> wrote:

@daverim https://github.com/daverim thank you for your reply. I tried your WAR and I'm still getting the same error:

ValueError: Unable to save function b'__inference_sep_conv2d_1_layer_call_fn_1381' because it captures graph tensor Tensor("new_model/toy_separableconv2d_test/quant_sep_conv2d_1/LastValueQuant_1/QuantizeAndDequantizeV4:0", shape=(3, 3, 1, 1), dtype=float32) from a parent function which cannot be converted to a constant with tf.get_static_value.

New code snippet:

model = quantize_model(model)model_save_path = 'weights/sample_qat' class NewModel(tf.keras.Model): def init(self, my_model): super(NewModel, self).init() self.my_model = my_model
def call(self, x):
    return self.my_model(x)
q_model = NewModel(q_model)q_model(tf.random.normal(shape=(1, 3, 224, 224)))q_model.save(model_save_path)

Please note that the code snippet I sent when opening this bug works for models with Conv2D, Dense, and many other layers, but fails when I include SeparableConv2D or ConvTranspose2D in the model. It also failed with DethwiseConv2D in tensorflow < 2.8. Can you think of any reason why? Tnank you!

— Reply to this email directly, view it on GitHub https://github.com/tensorflow/model-optimization/issues/964#issuecomment-1110386969, or unsubscribe https://github.com/notifications/unsubscribe-auth/AASV4JIER2YMPPMCLAPXPCLVHCE4DANCNFSM5TU2Q5BQ . You are receiving this because you were mentioned.Message ID: @.***>

gcunhase commented 2 years ago

Got the same error:

ValueError: Unable to save function b'__inference_sep_conv2d_1_layer_call_fn_1381' because it captures graph tensor Tensor("new_model/toy_separableconv2d_test/quant_sep_conv2d_1/LastValueQuant_1/QuantizeAndDequantizeV4:0", shape=(3, 3, 1, 1), dtype=float32) from a parent function which cannot be converted to a constant with `tf.get_static_value`.

New code snippet:

nn_model = _model()

class NewModel(tf.keras.Model):
    def __init__(self, my_model):
        super(NewModel, self).__init__()
        self.my_model = quantize_model(my_model)

    def call(self, x):
        return self.my_model(x)

q_model = NewModel(nn_model)
q_model(tf.random.normal(shape=(1, 3, 224, 224)))
q_model.save("saved_model_qat")

The bug's description contains the full repro.

Also, any comment on this? "The code snippet I sent when opening this bug works for models with Conv2D, Dense, and many other layers, but fails when I include SeparableConv2D or ConvTranspose2D in the model. It also failed with DethwiseConv2D in tensorflow < 2.8 (working in tf == 2.8). Can you think of any reason why?"

gcunhase commented 2 years ago

Any update on this?

gcunhase commented 2 years ago

@daverim I wonder how this very similar issue got solved: https://github.com/tensorflow/model-optimization/issues/868#issuecomment-1006225011

gcunhase commented 2 years ago

Any update?

gcunhase commented 2 years ago

Any update?

gcunhase commented 2 years ago

Any update? @daverim

gcunhase commented 2 years ago

Any update?

gcunhase commented 1 year ago

Any update?

gcunhase commented 1 year ago

Any update?

daverim commented 1 year ago

Also, i'm not sure what quantize_model is doing, but basically, there are some tensors that are outside the function def that you want to save. The code I presented encapsulates all the data inside a function.

If you are using QAT api then there is a transform that converts a separable conv to a conv2d and depthwise conv2d with a quantize_and_dequantize between the two -- if you are using this functionality outside of quantize_apply then this might not work correctly, so I highly recommend just using the qat api as directed in the examples.

On Wed, Apr 27, 2022 at 10:11 AM David Rim @.***> wrote:

you need to move quantize_model into the call or the init
self.my_model = quantize_model(model)
On Wed, Apr 27, 2022 at 9:45 AM Gwena Cunha @.***> wrote:
@daverim https://github.com/daverim thank you for your reply. I tried your WAR and I'm still getting the same error:

ValueError: Unable to save function b'__inference_sep_conv2d_1_layer_call_fn_1381' because it captures graph tensor Tensor("new_model/toy_separableconv2d_test/quant_sep_conv2d_1/LastValueQuant_1/QuantizeAndDequantizeV4:0", shape=(3, 3, 1, 1), dtype=float32) from a parent function which cannot be converted to a constant with tf.get_static_value.

New code snippet:

model = quantize_model(model)model_save_path = 'weights/sample_qat' class NewModel(tf.keras.Model): def init(self, my_model): super(NewModel, self).init() self.my_model = my_model
def call(self, x):
    return self.my_model(x)
q_model = NewModel(q_model)q_model(tf.random.normal(shape=(1, 3, 224, 224)))q_model.save(model_save_path)

Please note that the code snippet I sent when opening this bug works for models with Conv2D, Dense, and many other layers, but fails when I include SeparableConv2D or ConvTranspose2D in the model. It also failed with DethwiseConv2D in tensorflow < 2.8. Can you think of any reason why? Tnank you!

— Reply to this email directly, view it on GitHub https://github.com/tensorflow/model-optimization/issues/964#issuecomment-1110386969, or unsubscribe https://github.com/notifications/unsubscribe-auth/AASV4JIER2YMPPMCLAPXPCLVHCE4DANCNFSM5TU2Q5BQ . You are receiving this because you were mentioned.Message ID: @.***>

tensorflow / model-optimization

QAT model saving bug: Unable to save function b'__inference_separable_conv2d_layer_call_fn_961' #964