如何使用TF 2.0从巨大的二进制数据中生成训练数据集？

import tensorflow as tf
import numpy as np
import os

def data_generator(data_dir, batch_size):
    # 读取二进制数据文件
    data = np.fromfile(data_dir, dtype=np.float32)
    # 根据数据大小计算总的样本数
    num_samples = len(data) // input_size
    # 将数据重塑为合适的形状
    data = data.reshape((num_samples, input_size))
    # 创建一个无限循环的生成器
    while True:
        # 随机打乱数据
        np.random.shuffle(data)
        # 生成批量数据
        for i in range(0, num_samples, batch_size):
            batch_data = data[i:i+batch_size]
            yield batch_data

# 定义模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu', input_shape=(input_size,)),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(num_classes, activation='softmax')
])

# 编译模型
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# 创建数据集生成器
train_generator = data_generator(train_data_dir, batch_size)

# 训练模型
model.fit(train_generator,
          steps_per_epoch=train_steps_per_epoch,
          epochs=num_epochs)