def create_dataset_from_tf_record_files(input_file_pattern, pre_batch_size, batch_size, is_training=True):
    'Creates dataset from (tf)records files for training/evaluation.'
    files = tf.data.Dataset.list_files(input_file_pattern, shuffle=is_training)

    def make_dataset(files_dataset, shard_index):
        'Returns dataset for sharded tf record files.'
        if (pre_batch_size != batch_size):
            raise ValueError('Pre-batch ({}) size is not equal to batch size ({})'.format(pre_batch_size, batch_size))
        files_dataset = files_dataset.shard(NUM_SHARDS, shard_index)
        dataset = files_dataset.interleave(tf.data.TFRecordDataset)
        decode_fn = functools.partial(data_pipeline.DatasetManager.deserialize, batch_size=pre_batch_size, is_training=is_training)
        dataset = dataset.map(decode_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
        return dataset
    dataset = tf.data.Dataset.range(NUM_SHARDS)
    map_fn = functools.partial(make_dataset, files)
    dataset = dataset.interleave(map_fn, cycle_length=NUM_SHARDS, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return dataset
Creates dataset from (tf)records files for training/evaluation.
def create_dataset_from_tf_record_files(input_file_pattern, pre_batch_size, batch_size, is_training=True):
    
    files = tf.data.Dataset.list_files(input_file_pattern, shuffle=is_training)

    def make_dataset(files_dataset, shard_index):
        'Returns dataset for sharded tf record files.'
        if (pre_batch_size != batch_size):
            raise ValueError('Pre-batch ({}) size is not equal to batch size ({})'.format(pre_batch_size, batch_size))
        files_dataset = files_dataset.shard(NUM_SHARDS, shard_index)
        dataset = files_dataset.interleave(tf.data.TFRecordDataset)
        decode_fn = functools.partial(data_pipeline.DatasetManager.deserialize, batch_size=pre_batch_size, is_training=is_training)
        dataset = dataset.map(decode_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
        return dataset
    dataset = tf.data.Dataset.range(NUM_SHARDS)
    map_fn = functools.partial(make_dataset, files)
    dataset = dataset.interleave(map_fn, cycle_length=NUM_SHARDS, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return dataset