如何使用多个.npy 文件创建 tf.data 管道

Question

我已经研究了有关此问题的其他问题，但找不到确切答案，因此从头开始尝试：

问题

我有多个 .npy 文件（X_train 文件），每个文件都是一个形状数组（n、99、2）——只有第一个维度不同，其余两个相同。 根据 .npy 文件的名称，我还可以获得相应的标签（y_train 文件）。

每一对这样的文件都可以很容易地加载到 memory 中（多个文件也是如此），但不是一次全部加载。

我构建了一个生成器，它遍历文件列表并为训练批次聚合给定数量的文件：

def tf_data_generator(filelist, directory = [], batch_size = 5):
    i = 0
    x_t = os.listdir(directory[0])
    y_t = os.listdir(directory[1])
    while True:
        file_chunk = filelist[i*batch_size:(i+1)*batch_size] 
        X_a = []
        Y_a = []
        for fname in file_chunk:
            x_info = np.load(path_x_tr+fname)
            y_info = np.load(path_y_tr+fname)
            X_a.append(x_info)
            Y_a.append(y_info)
        X_a = np.concatenate(X_a)
        Y_a = np.concatenate(Y_a)
        yield X_a, Y_a
        i = i + 1

在实践中（在 CPU 上）它工作正常，但是如果我试图在 CUDA 上使用 GPU，它会崩溃，给出 Failed to call ThenRnnForward Failed to call ThenRnnForward with model config:

所以我试图找到另一种方法并使用 tf.data API 进行数据生成。 但是，我被困住了：

def parse_file(name):
    x = np.load('./data/x_train_m/'+name)
    y = np.load('./data/y_train_m/'+name)
    train_dataset = tf.data.Dataset.from_tensor_slices((test1, test2))
    return train_dataset

train_dataset = parse_file('example1.npy')
train_dataset = train_dataset.shuffle(100).batch(64)

model = wtte_rnn()
model.summary()
K.set_value(model.optimizer.lr, 0.01)
model.fit(train_dataset,
          epochs=10)

这很好用，但是，我找不到以下方法：

1. 混合多个文件（最多一定数量，比如说五个）
2. 遍历整个文件列表

我已经阅读了 flat_map 和 interleave，但是，我无法进一步使用 go 并且任何尝试使用它们都没有成功。 如何制作与代码上部类似的生成器，但使用 tf.data API？

Answer 1

您可以尝试连接它们，如下所示：

train_dataset = parse_file('example1.npy') # initialize train dataset

for file in files[1:]: # concatenate with the remaining files
    train_dataset = train_dataset.concatenate(parse_file(file))