Tensorflow Keras数据集加载大量Numpy文件

import tensorflow as tf
import numpy as np
import os

def load_numpy_file(file_path):
    data = np.load(file_path)
    # 假设数据和标签分别保存在'data'和'label'两个键中
    x = data['data']
    y = data['label']
    return x, y

def load_dataset(data_dir):
    file_paths = [os.path.join(data_dir, file) for file in os.listdir(data_dir) if file.endswith('.npy')]
    dataset = tf.data.Dataset.from_tensor_slices(file_paths)
    dataset = dataset.map(load_numpy_file)
    return dataset

data_dir = '/path/to/dataset'
dataset = load_dataset(data_dir)

# 划分训练集和测试集
train_dataset = dataset.take(800)
test_dataset = dataset.skip(800)

# 打乱数据
train_dataset = train_dataset.shuffle(800)

# 批量处理数据
batch_size = 32
train_dataset = train_dataset.batch(batch_size)
test_dataset = test_dataset.batch(batch_size)