如何在Tensorflow中实现多线程导入存储在磁盘上的numpy数组作为数据集

Question

我的数据集的输入和标签分别存储在 10000 个 .npy 文件中。例如 inputs/0000.npy,...inputs/9999.npy 和 labels/0000.npy,...labels/9999.npy。虽然每个文件都可以独立存储在内存中，但 20k 数组的整个数据集不能存储在内存中。我想实现多线程 CPU 管道以将数据集导入为 batch_size=8 的批次。

我已经尝试实现新的 Tensorflow 数据中提到的功能 API 但没有找到符合我要求的示例。所有示例似乎都是针对可以将整个数据集加载到 RAM 中的情况。知道如何解决这个问题吗？

Answer 1

我会使用 tf.data.Dataset.from_generator()，它允许您通过自定义 python 生成器函数使用 Tensorflow 数据 API。这样，您可以迭代地加载每个 .npy 文件，一次只在内存中加载一个 numpy.ndarray 文件。假设每个加载的 numpy.ndarray 都是一个实例，您的案例的示例代码可能如下所示：

import tensorflow as tf
import numpy as np
import os


def gen():
    inputs_path = ""
    labels_path = ""
    for input_file, label_file in zip(os.listdir(inputs_path), os.listdir(labels_path)):
        x = np.load(os.path.join(inputs_path, input_file))
        y = np.load(os.path.join(labels_path, label_file))
        yield x, y


INPUT_SHAPE = []
LABEL_SHAPE = []

# Input pipeline
ds = tf.data.Dataset.from_generator(
    gen, (tf.float32, tf.int64), (tf.TensorShape(INPUT_SHAPE), tf.TensorShape(LABEL_SHAPE)))
ds = ds.batch(8)
ds_iter = ds.make_initializable_iterator()
inputs_batch, labels_batch = ds_iter.get_next()

我没有测试代码。希望对您有所帮助！

如何在Tensorflow中实现多线程导入存储在磁盘上的numpy数组作为数据集

How to implement multi-threaded import of numpy arrays stored on disk as dataset in Tensorflow

python

tensorflow

tensorflow-datasets