tensorflow: Segmentation fault using tf.lite.TFLiteConverter with representative_dataset

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): No
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: No
• TensorFlow installed from (source or binary): Source
• TensorFlow version (use command below): 2.0beta1
• Python version: 3.6.8
• Bazel version (if compiling from source): 0.24.1
• GCC/Compiler version (if compiling from source): 7.4.0
• CUDA/cuDNN version: 10.0
• GPU model and memory: GTX2080TI 11 GB
• Exact command to reproduce:

Describe the problem

I am trying to quantize (post-training) my model, using a representative dataset as explained in the latest documentation. I tried to build the representative dataset by loading some relevant PNGs from a folder. But I get a segmentation fault error when I set the representative dataset (no error if I do not set it or use mnist as dataset).

Source code / logs

This works:

mnist_train, _ = tf.keras.datasets.mnist.load_data()
images = tf.cast(mnist_train[0], tf.float32)/255.0
mnist_ds = tf.data.Dataset.from_tensor_slices((images)).batch(1)
def representative_data_gen():
  for input_value in mnist_ds.take(100):
    yield [input_value]

tf.logging.set_verbosity(tf.logging.INFO)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()

This doesn’t (SEGFAULT):

train = []
files = glob.glob ('/myPath/data/images/test/*.png')

for myFile in files:
    image = cv2.imread (myFile)
    train.append (image)
train = np.array(train,dtype='float32') #as mnist
images = tf.cast(train[0], tf.float32)/255.0
my_ds = tf.data.Dataset.from_tensor_slices((images)).batch(1)
    
def representative_data_gen():
  for input_value in my_ds.take(100):
    yield [input_value]    

tf.logging.set_verbosity(tf.logging.INFO)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()

Output:

....
2019-06-15 20:40:32.261287: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:1006] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero
2019-06-15 20:40:32.261634: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1326] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 10246 MB memory) -> physical GPU (device: 0, name: GeForce RTX 2080 Ti, pci bus id: 0000:01:00.0, compute capability: 7.5)
W0615 20:40:32.319736 140096456644416 deprecation_wrapper.py:118] From quantize.py:47: The name tf.logging.INFO is deprecated. Please use tf.compat.v1.logging.INFO instead.

Fatal Python error: Segmentation fault

Current thread 0x00007f6abf885740 (most recent call first):
  File "/usr/tf2/lib/python3.6/site-packages/tensorflow/lite/python/optimize/calibrator.py", line 51 in __init__
  File "/usr/tf2/lib/python3.6/site-packages/tensorflow/lite/python/lite.py", line 197 in _calibrate_quantize_model
  File "/usr/tf2/lib/python3.6/site-packages/tensorflow/lite/python/lite.py", line 922 in convert
  File "quantize.py", line 50 in <module>

Segmentation fault (core dumped)

If I use ciphar instead of mnist as representative dataset I also get the same SEGFAULT.

images = tf.cast(cifar_train[0], tf.float32) / 255.0
cifar_ds = tf.data.Dataset.from_tensor_slices((images)).batch(1)
def representative_data_gen():
  for input_value in cifar_ds.take(100):
    yield [input_value]`

Any ideas? Thank you.