如何使用Keras读取大型parquet文件并传入多输入模型？

方案1：tf.data + TensorFlow I/O 原生读取（超大数据集首选） #

该方案无需将全量数据加载到内存，完全适配Keras训练流水线，是处理TB级数据集的最优选择

• 首先安装依赖：pip install tensorflow-io
• 代码实现示例：

import tensorflow as tf
import tensorflow_io as tfio

def parse_parquet_record(record):
    # 字段类型需和你存储parquet时的字段类型完全对应
    feature_description = {
        "input_col1": tf.io.FixedLenFeature([], tf.float32), # 单值输入
        "input_col2": tf.io.FixedLenFeature([128], tf.float32), # 长度128的向量输入
        "input_col3": tf.io.FixedLenFeature([10], tf.int32), # 长度10的离散特征输入
        "label_col": tf.io.FixedLenFeature([], tf.int32) # 标签列
    }
    parsed_data = tf.io.parse_single_example(record, feature_description)
    # 按你的Keras多输入定义返回对应结构，用字典或元组都可
    model_inputs = {
        "model_input1": parsed_data["input_col1"],
        "model_input2": parsed_data["input_col2"],
        "model_input3": parsed_data["input_col3"]
    }
    return model_inputs, parsed_data["label_col"]

# 支持传入单个或多个parquet文件路径列表
parquet_file_paths = ["data_part1.parquet", "data_part2.parquet"]
# 构造数据流水线
dataset = tf.data.Dataset.from_tensor_slices(parquet_file_paths)
# 并行读取多个文件
dataset = dataset.interleave(
    lambda path: tfio.IODataset.from_parquet(path),
    num_parallel_calls=tf.data.AUTOTUNE
)
# 解析每条记录
dataset = dataset.map(parse_parquet_record, num_parallel_calls=tf.data.AUTOTUNE)
# 标准训练流水线配置，buffer和batch_size可根据你的硬件调整
dataset = dataset.shuffle(buffer_size=20000).batch(64).prefetch(tf.data.AUTOTUNE)

# 直接传入Keras模型训练即可
model.fit(dataset, epochs=15)

方案2：pandas分块读取+生成器适配（中等体量数据集首选） #

该方案无需额外依赖，调试成本低，适合无法全量加载到内存但不需要分布式训练的场景

• 代码实现示例：

import pandas as pd
import numpy as np
import tensorflow as tf

# 每块读取的行数，根据你的可用内存调整
CHUNK_SIZE = 50000
BATCH_SIZE = 64

def parquet_data_generator(parquet_path):
    # 迭代分块读取parquet
    for chunk in pd.read_parquet(parquet_path, chunksize=CHUNK_SIZE):
        # 拆分多输入和标签
        input1 = chunk["input_col1"].values
        input2 = np.stack(chunk["input_col2"].values)
        input3 = np.stack(chunk["input_col3"].values)
        labels = chunk["label_col"].values
        # 按批次yield数据
        for idx in range(0, len(chunk), BATCH_SIZE):
            batch_in1 = input1[idx:idx+BATCH_SIZE]
            batch_in2 = input2[idx:idx+BATCH_SIZE]
            batch_in3 = input3[idx:idx+BATCH_SIZE]
            batch_label = labels[idx:idx+BATCH_SIZE]
            yield (batch_in1, batch_in2, batch_in3), batch_label

# 构造tf数据集，output_signature需要和你实际数据的形状、类型对应
dataset = tf.data.Dataset.from_generator(
    lambda: parquet_data_generator("full_data.parquet"),
    output_signature=(
        (
            tf.TensorSpec(shape=(None,), dtype=tf.float32),
            tf.TensorSpec(shape=(None,128), dtype=tf.float32),
            tf.TensorSpec(shape=(None,10), dtype=tf.int32)
        ),
        tf.TensorSpec(shape=(None,), dtype=tf.int32)
    )
).prefetch(tf.data.AUTOTUNE)

# 直接传入模型训练
model.fit(dataset, epochs=15)

注意事项 #

• 多输入返回的结构、字段顺序/名称必须和你定义Keras模型时的输入结构完全匹配
• 如果parquet中存在变长序列特征，将FixedLenFeature替换为VarLenFeature，解析后补充padding操作即可
• 分布式训练场景优先选择方案1，原生tf.data的兼容性和性能表现更优

内容的提问来源于stack exchange，提问作者Lostsoul