TFLite Converter segfaults when trying to convert per-channel quantized transposed convolutions #53767
Description
When converting transposed convolutions using per-channel weight quantization the converter segfaults and crashes the Python process. Per-channel quantization is supported by TFLite Transposed convolutions:
tensorflow/tensorflow/lite/kernels/transpose_conv.cc
Lines 371 to 380 in f87be6c
so the converter shouldn't segfault when trying to convert such a model.
It looks like this issue has been introduced in TensorFlow 2.6 since the same model code produced a valid TFLite file in TensorFlow 2.5. This issue might also be related to #53766, but in any case the converter should never segfault.
1. System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS / Ubuntu
- TensorFlow installation (pip package or built from source): pip package
- TensorFlow library (version, if pip package or github SHA, if built from source): 2.6, 2.7, 2.8rc0 and 2.9.0-dev20220114
2. Code
A minimal reproduction of the issue and a workaround is available in this notebook.
import tensorflow as tf
class QuantConv2DTransposed(tf.keras.layers.Layer):
def build(self, input_shape):
self.kernel = self.add_weight("kernel", [3, 3, input_shape[-1], 24])
def call(self, inputs):
filters = tf.quantization.fake_quant_with_min_max_vars_per_channel(
self.kernel, -3.0 * tf.ones([24]), 3.0 * tf.ones([24]), narrow_range=True
)
filters = tf.transpose(filters, (0, 1, 3, 2))
return tf.nn.conv2d_transpose(inputs, filters, [*inputs.shape[:-1], 24], 1)
inp = tf.keras.Input(shape=(6, 8, 48), batch_size=1)
x = tf.quantization.fake_quant_with_min_max_vars(inp, -3.0, 3.0, narrow_range=True)
x = QuantConv2DTransposed()(x)
x = tf.quantization.fake_quant_with_min_max_vars(x, -3.0, 3.0, narrow_range=True)
model = tf.keras.Model(inp, x)
model.save("/tmp/testing")
converter = tf.lite.TFLiteConverter.from_saved_model("/tmp/testing")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# terminated by signal SIGSEGV (Address boundary error)
tflite_model = converter.convert()