Filtering batches from multiple files in TensorFlow

Question

I'm reading TF examples from multiple files and want to filter the examples before passing them off to some batched operations.

Unfortunately, both approaches I've considered for this seem to fail.

An MWE follows.

Create dataset and define a couple of useful functions:

import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow.compat.v1 as tfv1
tfv1.enable_v2_behavior()

def _float_feature(value):
  """Returns a float_list from a float / double."""
  return tf.train.Feature(float_list=tf.train.FloatList(value=value))

def serialize_example(feature0, feature1):
  """
  Creates a tf.Example message ready to be written to a file.
  """
  # Create a dictionary mapping the feature name to the tf.Example-compatible
  # data type.
  feature = {
      'filter_on_this': _float_feature(feature0),
      'sequence':       _float_feature(feature1),
  }

  example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
  return example_proto.SerializeToString()

for i in range(10):
  # Write the `tf.Example` observations to the file.
  with tf.io.TFRecordWriter("out{0}.tfexamples".format(i)) as writer:
    for i in range(1000):
      example = serialize_example(np.random.random(1), np.random.random(np.random.randint(50,70)))
      writer.write(example)

sequence_numeric_features_spec = {
    'filter_on_this': tf.io.FixedLenFeature((1,), dtype=tf.float32, default_value=0), #tf.io.VarLenFeature(dtype=tf.float32)
    'sequence':       tf.io.VarLenFeature(dtype=tf.float32),
}

DEFAULT_DTYPE_VALUES = {
    tf.float32: 0.,
    tf.string: '',
    tf.int64: 0,
}

def parse_record2(record):
  record = tf.io.parse_example(record, features=sequence_numeric_features_spec)
  for key, feature in record.items():
    if not isinstance(feature, tf.SparseTensor):
      continue
    dense = tf.sparse.to_dense(sp_input=feature, default_value=DEFAULT_DTYPE_VALUES[feature.values.dtype])
    dense = tf.expand_dims(dense, axis=2)
    record[key] = dense
  return record

First approach:

files   = tf.data.Dataset.list_files(file_pattern="out*.tfexamples", shuffle=True, seed=123456789)
dataset = files.interleave(lambda x: tf.data.TFRecordDataset(x).prefetch(100), cycle_length=8, num_parallel_calls=4)
dataset = dataset.batch(10)
dataset = dataset.map(parse_record2, num_parallel_calls=20)
dataset = dataset.filter(lambda x: x['filter_on_this']<0.5)
for x in dataset.take(1):
  pass

This approach is desirable because when examples are parsed independently the decoding conversion in parse_record2 takes a long time. Batching the examples speeds up the conversion quite a bit. Unfortunately, this raises the error:

ValueError: predicate return type must be convertible to a scalar boolean tensor. Was TensorSpec(shape=(None, 1), dtype=tf.bool, name=None).

Second approach

Since batches seem to be problematic, I've tried an alternative without them:

files   = tf.data.Dataset.list_files(file_pattern="out*.tfexamples", shuffle=True, seed=123456789)
dataset = files.interleave(lambda x: tf.data.TFRecordDataset(x).prefetch(100), cycle_length=8, num_parallel_calls=4)
dataset = dataset.batch(20)
dataset = dataset.map(parse_record2, num_parallel_calls=20)
dataset = dataset.unbatch()
dataset = dataset.filter(lambda x: x['filter_on_this'][0]<0.5)
dataset = dataset.batch(20)
for x in dataset.take(4):
  pass

But, in addition to seeming like a less performant approach, this fails with

InvalidArgumentError: Cannot batch tensors with different shapes in component 1. First element had shape [69,1] and element 2 had shape [66,1].

What's a good way of decoding examples in chunks and filtering them into equally-sized batches?

Filtering batches from multiple files in TensorFlow

0 Answers