QAT model saving bug: Unable to save function b'__inference_separable_conv2d_layer_call_fn_961'
gcunhase opened this issue · 15 comments
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
DepthwiseConv2Dintf<2.8.0, but has been fixed in 2.8.0. Please see the previous issue here: #868 . - This issue also happens with
ConvTranspose2D.
Any updates?
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.
@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!
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?"
Any update on this?
@daverim I wonder how this very similar issue got solved: #868 (comment)
Any update?
Any update?
Any update?
Any update?
Any update?