TensorFlow - 交错多个独立预处理的 TFRecord 文件

Question

I have multiple TFRecord files from the Waymo Dataset, each containing consecutive points that are not consecutive across files. I'm building an input pipeline that preprocesses data for time series prediction via the window() API but I need to avoid the window to span accross multiple files.

To do so, I believe I should preprocess each file indipentently and interleave the final datasets. Here's my attempt:

import tensorflow as tf
from waymo_open_dataset import dataset_pb2 as open_dataset #for parsing Waymo frames

filenames = [os.path.join(DATASET_DIR, f) for f in os.listdir(DATASET_DIR)]
dataset = tf.data.TFRecordDataset(filenames, compression_type='')

def interleave_fn(filename):
    ds = filename.map(lambda x: tf.py_function(_parse_data, [x], [tf.float32]*N_FEATURES,), 
                          num_parallel_calls=tf.data.experimental.AUTOTUNE) 
    ds = ds.map(_concatenate_tensors).map(_set_x_shape)
    ds = build_x_dataset(ds)
    return ds

def _parse_data(data):
    # Parse feature from Waymo dataset  
    frame = open_dataset.Frame()
    frame.ParseFromString(bytearray(data.numpy()))   
    av_v_x = frame.images[0].velocity.v_x 
    av_v_y = frame.images[0].velocity.v_y 
    return av_v_x, av_v_y

def _concatenate_tensors(*x):
    #Concatenate tensor tuple in a single tensor
    return tf.stack((x))

def _set_x_shape(x):
    #Set X dataset shape. If not UNDEFINED RANK ValueError
    x.set_shape((N_FEATURES,))
    return x
    
def build_x_dataset(ds_x, window = WINDOW):
    # Extract sequences for time series prediction training
    # Selects a sliding window of WINDOW samples, shifting by 1 sample at a time
    ds_x = ds_x.window(size = window, shift = 1, drop_remainder = True)
    
    # Each element of `ds_x` is a nested dataset containing WINDOWconsecutive examples 
    ds_x = ds_x.map(lambda d: tf.data.experimental.get_single_element(d.batch(window))) 
    return ds_x

dataset = dataset.interleave(interleave_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)

This returns

AttributeError: in user code:

    /tmp/xpython_26752/494049692.py:118 interleave_fn  *
        ds = filename.map(lambda x: tf.py_function(_parse_data, [x], [tf.float32]*N_FEATURES,),

    AttributeError: 'Tensor' object has no attribute 'map'

which makes sense because print(filename) in interleave_fn gives

Tensor("args_0:0", shape=(), dtype=string)

I thought the interleave_fn would be applied to each TFRecordDataset , so filename would be a dataset itself instead of a tensor. What's wrong here? Thank you!

Answer 1

Solved it by looping over all TFRecord files and appending the corresponding datasets to a dataset list. Then, following this tip to interleave all the preprocessed datasets.