如何在TensorFlow中使用tf.data读取.csv文件替代pandas处理大体积数据

tf.data支持流式读取逻辑，无需将全量数据载入内存，完全适配你多CSV按列拼接、喂入模型训练的需求，具体实现如下：

1. 前置注意事项 #

请保证三份CSV的样本顺序完全一一对应，才能按你原逻辑横向拼接特征，和pd.concat(axis=1)效果一致。

2. 完整实现代码 #

import tensorflow as tf

# 替换为你的三份CSV真实路径
CSV_PATHS = ["data1.csv", "data2.csv", "data3.csv"]
# 每份CSV的特征列数，你当前场景为3，根据实际调整
PER_FILE_FEATURES = 3
# 训练批次大小，根据显存容量调整
BATCH_SIZE = 32

def load_single_csv(csv_path):
    # 流式读取CSV，跳过表头行
    dataset = tf.data.TextLineDataset(csv_path).skip(1)
    # 解析每行内容为浮点型特征
    def _parse_line(line):
        fields = tf.io.decode_csv(line, record_defaults=[[0.0] for _ in range(PER_FILE_FEATURES)])
        return tf.stack(fields, axis=0)
    return dataset.map(_parse_line, num_parallel_calls=tf.data.experimental.AUTOTUNE)

# 并行加载三份CSV
ds1 = load_single_csv(CSV_PATHS[0])
ds2 = load_single_csv(CSV_PATHS[1])
ds3 = load_single_csv(CSV_PATHS[2])

# 按样本对齐，横向拼接三份特征
combined_ds = tf.data.Dataset.zip((ds1, ds2, ds3))\
    .map(lambda x,y,z: tf.concat([x,y,z], axis=0), num_parallel_calls=tf.data.experimental.AUTOTUNE)

# 配置打乱、分批、预取逻辑，优化读写性能
combined_ds = combined_ds.shuffle(buffer_size=1000)\
    .batch(BATCH_SIZE)\
    .prefetch(tf.data.experimental.AUTOTUNE)

# 适配tf1.x会话的迭代器定义
iterator = tf.compat.v1.data.make_initializable_iterator(combined_ds)
next_batch = iterator.get_next()

替换你原有会话部分的训练逻辑即可：

# 原有模型定义、优化器定义、init初始化操作保持不变
with tf.compat.v1.Session() as sess:
    sess.run(init)
    # 初始化数据集迭代器
    sess.run(iterator.initializer)
    # 迭代训练
    while True:
        try:
            batch_data = sess.run(next_batch)
            sess.run(optimizer, feed_dict={SOME_VARIABLE: batch_data})
        except tf.errors.OutOfRangeError:
            # 数据集遍历完一轮，可重启迭代器继续下一轮训练，或直接退出
            break

3. 调优提示 #

• 若你的CSV包含索引列，调整decode_csv的参数，解析后跳过索引列即可
• shuffle的buffer_size可根据可用内存调整，数值越大打乱效果越好，无需超过总样本数
• 若需要多轮训练，每次触发OutOfRangeError后重新执行sess.run(iterator.initializer)即可重新加载数据集

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