如何将特征字典构成的FlatMapDataset高效转换为单个大Tensor？

搭建了从TFRecord加载数据的流水线，预处理后得到FlatMapDataset——该数据集包含每条样本的特征字典（整体维度为nrow×features）。希望将其转换为一个大张量，当前仅能输出MapDataset，且遍历每个元素的方法在大数据集下效率过低。

当前代码：

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

# 下载数据集
!wget -q https://storage.googleapis.com/download.tensorflow.org/data/palmer_penguins/penguins.csv -O /tmp/penguins.csv

# 加载到Pandas DataFrame
dataset_df = pd.read_csv("/tmp/penguins.csv")

# 特征变量
predictors = ["bill_length_mm", "bill_depth_mm", "flipper_length_mm", "body_mass_g"]

# 筛选感兴趣的特征
dataset_df = dataset_df.loc[:, predictors]

# 创建FlatMapDataset
ds = tfdf.keras.pd_dataframe_to_tf_dataset(dataset_df)
ds = ds.flat_map(lambda x: tf.data.Dataset.from_tensor_slices(x))

# 展平
ds_flat = ds.map(lambda elem: tf.convert_to_tensor(tf.nest.flatten(elem)))

低效的遍历方法：

# 备选方案（效率低）
list_of_tensors = []
for element in ds:
    tensor = tf.convert_to_tensor(tf.nest.flatten(element))
    list_of_tensors.append(tensor)
ds_as_tensor = tf.convert_to_tensor(list_of_tensors)

期望输出：

print(ds_as_tensor)

<tf.Tensor: shape=(344, 4), dtype=float64, numpy=
array([[  18.7,   39.1, 3750. ,  181. ],
       [  17.4,   39.5, 3800. ,  186. ],
       [  18. ,   40.3, 3250. ,  195. ],
       ...,
       [  18.2,   49.6, 3775. ,  193. ],
       [  19. ,   50.8, 4100. ,  210. ],
       [  18.7,   50.2, 3775. ,  198. ]])>

方法1：跳过不必要的FlatMap操作（最优） #

tfdf.keras.pd_dataframe_to_tf_dataset默认返回的Dataset仅包含单个元素，该元素是一个字典，每个键对应特征列的完整张量（形状为(n_samples,)）。直接提取这些特征并拼接即可得到大张量，无需拆分后再合并：

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

# 下载数据集
!wget -q https://storage.googleapis.com/download.tensorflow.org/data/palmer_penguins/penguins.csv -O /tmp/penguins.csv

# 加载到Pandas DataFrame并处理缺失值
dataset_df = pd.read_csv("/tmp/penguins.csv")
predictors = ["bill_length_mm", "bill_depth_mm", "flipper_length_mm", "body_mass_g"]
dataset_df = dataset_df.loc[:, predictors].dropna()  # 移除含缺失值的样本

# 转换为TF Dataset（默认batch_size=None，包含整个数据集的特征字典）
ds = tfdf.keras.pd_dataframe_to_tf_dataset(dataset_df)

# 提取特征并拼接成二维张量
for full_batch in ds:
    # 按特征顺序提取张量，堆叠成(n_samples, n_features)形状
    feature_tensors = [full_batch[col] for col in predictors]
    ds_as_tensor = tf.stack(feature_tensors, axis=1)
    break  # 仅需处理单个batch

print(ds_as_tensor)

方法2：针对已有的FlatMapDataset优化 #

如果已经得到了拆分后的FlatMapDataset，可通过batch()方法将所有样本合并为一个批量，再转换为张量，避免逐个遍历：

# 假设已得到FlatMapDataset `ds`
# 获取样本总数（若已知可直接传入，无需计算）
num_samples = tf.data.experimental.cardinality(ds).numpy()

# 将所有样本合并为一个batch
full_batch_ds = ds.batch(num_samples)

# 转换为大张量
for batch in full_batch_ds:
    feature_tensors = [batch[col] for col in predictors]
    ds_as_tensor = tf.stack(feature_tensors, axis=1)
    break

print(ds_as_tensor)

关键说明 #

• 两种方法均利用TensorFlow的批量操作，避免Python层面的逐个遍历，效率远高于手动收集列表再转换。
• 注意内存限制：若数据集过大无法一次性存入内存，此方法不适用（此时应保持Dataset流式处理）。

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