如何高效将含嵌套列表的大型Pandas DataFrame转为TensorFlow Dataset？

针对你150万条带高维嵌入的数据集，直接用from_tensor_slices或转NumPy导致内存爆炸的问题，核心原因是嵌套列表格式的Pandas列会把所有数据一次性加载到内存——150万条嵌入数据的总内存开销接近10GB，远超常规内存承载能力。以下是几个实用的优化方案：

方案一：分块读取+TF Dataset生成器 #

利用Pandas分块读取功能，每次只加载部分数据到内存，通过生成器流式喂给TF Dataset，避免一次性加载全量数据：

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

def data_generator(chunk_size=10000):
    # 替换为你的原始数据路径（csv/parquet均可）
    for chunk in pd.read_csv('your_data.csv', chunksize=chunk_size):
        # 单块数据转numpy数组，内存可控
        target_emb = np.array(chunk['target_product_embeddings'].tolist(), dtype=np.float32)
        hist_emb = np.array(chunk['hist_product_embeddings'].tolist(), dtype=np.float32)
        target = chunk['target'].values.astype(np.int32)
        
        # 按批次yield更高效
        for batch_idx in range(0, len(chunk), 64):
            batch_t_emb = target_emb[batch_idx:batch_idx+64]
            batch_h_emb = hist_emb[batch_idx:batch_idx+64]
            batch_y = target[batch_idx:batch_idx+64]
            yield ({'target_product_embeddings': batch_t_emb, 'hist_product_embeddings': batch_h_emb}, batch_y)

# 创建TF Dataset并指定输出格式
train_dataset = tf.data.Dataset.from_generator(
    data_generator,
    output_signature=(
        {
            'target_product_embeddings': tf.TensorSpec(shape=(None, 1536), dtype=tf.float32),
            'hist_product_embeddings': tf.TensorSpec(shape=(None, 10, 1536), dtype=tf.float32)
        },
        tf.TensorSpec(shape=(None,), dtype=tf.int32)
    )
)

# 数据流水线优化：预取+缓存（可选）
train_dataset = train_dataset.prefetch(tf.data.AUTOTUNE)

方案二：转Parquet格式后直接用TF读取 #

Parquet是高效的列式存储格式，原生支持嵌套数组结构，TF可以直接流式读取，无需先加载到Pandas：

第一步：转存为Parquet（仅需执行一次） #

# 若当前能加载小样本测试，直接转存
df.to_parquet('your_data.parquet', engine='pyarrow')
# 若原始数据太大无法全量加载，用分块读取+分块写入：
# chunk_list = []
# for chunk in pd.read_csv('raw_data.csv', chunksize=10000):
#     chunk_list.append(chunk)
# pd.concat(chunk_list).to_parquet('your_data.parquet', engine='pyarrow')

第二步：TF直接读取Parquet #

import tensorflow as tf

train_dataset = tf.data.Dataset.from_parquet(
    'your_data.parquet',
    columns=['target_product_embeddings', 'hist_product_embeddings', 'target'],
    output_signature=(
        {
            'target_product_embeddings': tf.TensorSpec(shape=(1536,), dtype=tf.float32),
            'hist_product_embeddings': tf.TensorSpec(shape=(10, 1536), dtype=tf.float32)
        },
        tf.TensorSpec(shape=(), dtype=tf.int32)
    )
)

# 批量处理+预取
train_dataset = train_dataset.batch(64).prefetch(tf.data.AUTOTUNE)

方案三：拆分嵌入向量为独立NumPy文件 #

把高维嵌入向量单独存为.npy文件，目标变量存为小文件，用TF Dataset直接读取，完全规避Pandas内存开销：

第一步：拆分保存（仅需执行一次） #

import numpy as np
import pandas as pd

# 若能全量加载数据，直接拆分
np.save('target_product_embeddings.npy', np.array(df['target_product_embeddings'].tolist(), dtype=np.float32))
np.save('hist_product_embeddings.npy', np.array(df['hist_product_embeddings'].tolist(), dtype=np.float32))
df['target'].to_csv('target.csv', index=False)

第二步：TF读取 #

import tensorflow as tf
import pandas as pd

# 用内存映射模式读取npy，避免全量加载
target_emb_ds = tf.data.Dataset.from_tensor_slices(np.load('target_product_embeddings.npy', mmap_mode='r'))
hist_emb_ds = tf.data.Dataset.from_tensor_slices(np.load('hist_product_embeddings.npy', mmap_mode='r'))
target_ds = tf.data.Dataset.from_tensor_slices(pd.read_csv('target.csv')['target'].values)

# 合并数据集
train_dataset = tf.data.Dataset.zip((
    {'target_product_embeddings': target_emb_ds, 'hist_product_embeddings': hist_emb_ds},
    target_ds
))

# 流水线优化
train_dataset = train_dataset.batch(64).prefetch(tf.data.AUTOTUNE)

通用优化提示 #

• 优先选择Parquet格式，兼顾存储效率和读取速度
• 用prefetch(tf.data.AUTOTUNE)让数据加载和模型训练并行，提升整体效率
• 根据内存情况调整batch size，避免单批次内存过载

内容的提问来源于stack exchange，提问作者Jurgita-ds