你能交错来自多个文件的 tf.data 数据集吗？

Posted 2023-02-16

【中文标题】你能交错来自多个文件的 tf.data 数据集吗？

【英文标题】：Can you interleave a tf.data dataset from multiple files?

【发布时间】：2020-12-01 14:13:33

【问题描述】：

我目前有一个数据集，该数据集分为三种不同的 npy 文件类型：一种包含所有浮点数的 x var，一种包含所有整数的 x var，另一种包含所有 y 标签。

为了遍历所有文件，我从像这样的生成器创建了一个 tf.data 数据集

def tf_data_generator(filelist_float, filelist_int, filelist_y, batch_size=4096):
    i=0
    while True:
        if i == len(filelist_y):
            i = 0
            indicies = np.arange(len(filelist_float))
            np.random.shuffle(indicies)
            filelist_float
        else:
            xf = np.load(filelist_float[i])
            xi = np.load(filelist_int[i])
            y = np.load(filelist_y[i])
            filefinish = False
            fileIndex = 0
            # In case files have different lengths
            fileEnd = np.min([xf.shape[0], xi.shape[0], y.shape[0]])
            while not filefinish:
                if fileIndex + batch_size >= fileEnd:
                    yield np.concatenate((xf[fileIndex:fileEnd], xi[fileIndex:fileEnd]), axis=1), y[fileIndex:fileEnd]
                    filefinish = True
                else:
                    yield np.concatenate((xf[fileIndex:fileIndex+batch_size], xi[fileIndex:fileIndex+batch_size]), axis=1), y[fileIndex:fileIndex+batch_size]
                    fileIndex += batch_size
            i+=1

 training_generator = tf.data.Dataset.from_generator(tf_data_generator, args = training_files_dir, output_types = (tf.float16, tf.int16)) 

我的问题是如果交错这个过程是否会更快，如果是这样，我可以用这三个单独的文件来做到这一点，还是必须将它们全部附加到一个文件中？需要注意的一点是，每个 npy 文件都可以加载到内存中，但是会有超过 10,000 个 npy 文件，所以我无法将所有文件都加载到内存中。

提前感谢您的任何想法！

【参考方案1】：

事实证明，是的，你绝对可以。如果有人好奇或卡住，这里是我的 TF2

在 call 函数中，生成器只接收 X_float_input 文件之一，并使用它来查找相应的 X_int_input 和 Y_input 文件。

class generator:
    def __call__(self, file, batch_size): #X_float_input.npy
        actual_file = file.decode("utf-8") 
        xf = np.load(actual_file)
        xi = np.load(actual_file.replace("float", "int"))
        y = np.load(actual_file.replace("X_float", "Y"))

        filefinish = False
        fileIndex = 0
        # In case files have different lengths
        fileEnd = np.min([xf.shape[0], xi.shape[0], y.shape[0]])
        while not filefinish:
            if fileIndex + batch_size >= fileEnd:
                yield np.concatenate((xf[fileIndex:fileEnd], xi[fileIndex:fileEnd]), axis=1), y[fileIndex:fileEnd]
                filefinish = True
            else:
                yield np.concatenate((xf[fileIndex:fileIndex+batch_size], xi[fileIndex:fileIndex+batch_size]), axis=1), y[fileIndex:fileIndex+batch_size]
                fileIndex += batch_size

# training_files_dir = ["X_float_input_1.npy", "X_float_input_2.npy"]
training_generator = tf.data.Dataset.from_tensor_slices(training_files_dir)
training_generator = training_generator.interleave(lambda filename: tf.data.Dataset.from_generator(
            generator(), 
            output_types=(tf.float16, tf.int16),
            args=(filename, BATCH_SIZE,)), num_parallel_calls=tf.data.experimental.AUTOTUNE).prefetch(tf.data.experimental.AUTOTUNE)
training_generator = training_generator.shuffle(10000, reshuffle_each_iteration=True).repeat(EPOCHS)#.prefetch(tf.data.experimental.AUTOTUNE).cache()

但需要注意的一点是，您需要使用 epoch 数调用重复函数，否则生成器将不会更新。