如何将含6000万行的Pandas DataFrame转换为TensorFlow数据集并实现保存与跨程序加载？

我来帮你搞定这个问题——6000万行的单列DataFrame转TF Dataset并跨程序使用，确实要注意大数据量下的内存和兼容性问题。你之前用pandas_tfrecords没生成文件，大概率是参数配置或者工具本身的局限性，我推荐用TensorFlow原生方法，更可靠且适配大规模数据：

一、正确转换并保存TensorFlow数据集 #

首先处理你的DataFrame，然后转成TF Dataset再保存，全程用原生TF工具，避免第三方库的坑：

步骤1：读取并预处理DataFrame #

先把CSV读进来，处理空值（因为你的数据里有句子或空值，空值会影响后续序列化）：

import pandas as pd
import tensorflow as tf

# 读取数据，指定字符串类型减少内存占用（替换成你的列名）
df = pd.read_csv('data.csv', dtype={'text_column': str})
# 把空值替换成空字符串，避免序列化报错
df['text_column'] = df['text_column'].fillna('')

步骤2：转换为TensorFlow Dataset #

把DataFrame的列转成TF Dataset，这里有两种方式，根据你的内存情况选：

方式A：内存足够时直接转换

# 从DataFrame列生成Dataset
dataset = tf.data.Dataset.from_tensor_slices(df['text_column'].values)

方式B：内存紧张时用生成器（适合超大数据量）

def data_generator():
    for text in df['text_column'].values:
        yield text

dataset = tf.data.Dataset.from_generator(
    data_generator,
    output_signature=tf.TensorSpec(shape=(), dtype=tf.string)
)

步骤3：保存数据集 #

推荐两种保存方式，按需选择：

方式1：保存为分片TFRecord文件（推荐大数据量）

把数据分成多个小文件，避免单个文件过大，方便后续并行加载：

# 分成10个分片，可根据你的数据大小调整
num_shards = 10
save_dir = './tfrecords'
# 创建保存目录
import os
os.makedirs(save_dir, exist_ok=True)

# 序列化函数：把文本转成TFRecord需要的Example格式
def serialize_text(text):
    feature = {
        'text': tf.train.Feature(bytes_list=tf.train.BytesList(value=[text.numpy()])),
    }
    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
    return example_proto.SerializeToString()

# 对每个元素序列化
serialized_dataset = dataset.map(lambda x: tf.py_function(serialize_text, [x], tf.string))

# 分片保存
for i in range(num_shards):
    shard_dataset = serialized_dataset.shard(num_shards=num_shards, index=i)
    writer = tf.io.TFRecordWriter(f'{save_dir}/data_shard_{i}.tfrecord')
    for serialized_example in shard_dataset:
        writer.write(serialized_example.numpy())
    writer.close()

方式2：直接保存Dataset（更简单）

用TF内置的保存方法，无需手动序列化，加载也更方便：

save_dir = './saved_tf_dataset'
# 保存Dataset，同时记录元素类型（方便后续加载）
element_spec = dataset.element_spec
tf.data.experimental.save(dataset, save_dir, element_spec=element_spec)

二、在另一个程序中加载数据集 #

根据你选择的保存方式，对应不同的加载方法：

加载分片TFRecord文件 #

import tensorflow as tf

# 反序列化函数：把TFRecord的Example转回文本
def deserialize_example(serialized_example):
    feature_desc = {
        'text': tf.io.FixedLenFeature([], tf.string),
    }
    example = tf.io.parse_single_example(serialized_example, feature_desc)
    return example['text']

# 加载所有分片文件
file_pattern = './tfrecords/data_shard_*.tfrecord'
files = tf.data.Dataset.list_files(file_pattern, shuffle=False)
# 并行读取多个分片
dataset = files.interleave(
    lambda x: tf.data.TFRecordDataset(x),
    num_parallel_calls=tf.data.AUTOTUNE
)
# 反序列化
dataset = dataset.map(deserialize_example, num_parallel_calls=tf.data.AUTOTUNE)

# 可选：添加批量、预取等优化（根据你的任务需求）
dataset = dataset.batch(32).prefetch(tf.data.AUTOTUNE)

加载直接保存的Dataset #

这种方式最省心，一行代码搞定：

import tensorflow as tf

dataset = tf.data.experimental.load('./saved_tf_dataset')
# 同样可以添加批量、预取等优化
dataset = dataset.batch(32).prefetch(tf.data.AUTOTUNE)

三、大数据量场景的额外优化建议 #

• 跳过Pandas直接读CSV：如果6000万行让内存压力很大，可以直接用TF读取CSV，不用先转Pandas，更高效：

dataset = tf.data.experimental.make_csv_dataset(
    'data.csv',
    batch_size=32,
    column_names=['text_column'],  # 替换成你的列名
    column_defaults=[tf.string],
    shuffle=False,
    num_epochs=1
)
# 提取文本列（因为make_csv_dataset返回的是字典格式）
dataset = dataset.map(lambda x: x['text_column'])

• 使用预取和并行处理：加载时用num_parallel_calls=tf.data.AUTOTUNE和prefetch，让数据处理和模型训练并行，提升整体速度。

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