如何拆分tensorflow数据集

我有一组存储在csv文件中的数据，目前我将其读出并存储在np中，然后将其传输到数据集，使用下面的代码

def read_data():
    with open(fname, "r") as f:
        lines = f.read().split("\n")
        header = lines[0].replace('"', "").split(',')
        lines = lines[1:]

        print(header)
        print(len(lines))

    float_data = np.zeros((len(lines), len(header) - 1))
    for i, line in enumerate(lines):
        values = [float(x) for x in line.split(",")[1:]]
        float_data[i, :] = values

    return tf.data.Dataset.from_tensor_slices(float_data)

然后我想定义一个生成器函数来从这个数据集获取数据，但是看起来数据集是不可下标的，比如numpy，我可以使用[:2]从中获取数据，但是数据集不能

我怎么做

下面是我使用numpy作为输入时的生成器函数（第一个参数是numpy）

我不确定数据集是否可以像我使用numpy时做的那样做

我可以在这个生成器的末尾使用

tf.data.Dataset.from_tensor_slices

方法，但是即使我使用

from_generator（generator.prefetch（）

方法，性能也很低，我假设这是因为数据非常大，性能受到CPU处理numpy数据的限制（我提到了这个问题），所以我想从一开始就以张量的形式加载数据，看看这是否会加快代码的速度

---

谢谢

您可以使用下面的示例代码拆分Tensorflow数据集

test_ds_size = int(length * 0.2) # 20 percent of length of ds
train_ds = ds.skip(test_ds_size)
test_ds = ds.take(test_ds_size)  

最好使用

预取

在que中加载一批进行训练，以提高训练速度。您可以按以下方式使用

预取

ds.batch(batch_size).prefetch(1)  

您也可以使用GPU和多个CPU来提高性能，下图所示

ds.batch(batch_size).prefetch(1)