如何使用多個.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))