使用TensorFlow Dataset API和Pandas构建输入管道的形状异常求助

看起来你遇到的核心问题是每个CSV文件的样本被当成了数据集的一个独立元素，而不是把所有文件的样本合并成一个扁平的数据集。让我一步步帮你分析和解决：

问题根源 #

当你用map处理每个文件路径时，你的解析函数返回的是一个(100,2)的张量（每个文件前100行，2列）。此时tf.data.Dataset的结构是：包含2个元素，每个元素是(100,2)的张量，所以整体呈现(2,100,2)的形状。而你需要的是把这2个元素里的所有样本展开，变成200个(2,)的样本，最终批次形状为(200,2)。

解决方案：用flat_map替代map（推荐） #

flat_map的作用是对每个输入元素（文件路径）生成一个子数据集，然后自动把所有子数据集拼接成一个大的扁平数据集。这样就能直接得到所有样本的集合，而不是按文件分组的张量。

修改后的代码示例 #

import tensorflow as tf
import pandas as pd

def parse_file(file_path):
    # 把TensorFlow的字符串张量转成Python字符串
    file_path_str = file_path.numpy().decode('utf-8')
    # 读取每个文件的前100行
    df = pd.read_csv(file_path_str, nrows=100)
    # 将当前文件的样本转换成一个小数据集（每个元素是(2,)的样本）
    return tf.data.Dataset.from_tensor_slices(
        tf.convert_to_tensor(df.values, dtype=tf.float32)
    )

# 你的CSV文件路径列表
file_paths = ["file1.csv", "file2.csv"]
# 创建初始数据集
dataset = tf.data.Dataset.from_tensor_slices(file_paths)

# 使用flat_map代替map，展开所有子数据集
# 注意：tf.py_function需要指定返回的数据集类型
dataset = dataset.flat_map(
    lambda x: tf.py_function(
        parse_file, 
        [x], 
        tf.data.DatasetSpec(tf.float32, shape=(2,))
    )
)

# 验证结果：生成200个样本的批次
for batch in dataset.batch(200):
    print(batch.shape)  # 输出 (200, 2)，符合预期

备选方案：map后加unbatch #

如果你不想修改解析函数，也可以在map之后用unbatch()把每个文件的(100,2)张量拆分成100个(2,)的样本，再合并：

def parse_file(file_path):
    file_path_str = file_path.numpy().decode('utf-8')
    df = pd.read_csv(file_path_str, nrows=100)
    return tf.convert_to_tensor(df.values, dtype=tf.float32)

dataset = tf.data.Dataset.from_tensor_slices(file_paths)
dataset = dataset.map(lambda x: tf.py_function(parse_file, [x], tf.float32))
dataset = dataset.unbatch()  # 拆分每个文件的张量为单个样本
dataset = dataset.batch(200)  # 生成目标批次

for batch in dataset:
    print(batch.shape)  # 同样输出 (200, 2)

关键注意事项 #

• 使用tf.py_function时，因为我们调用了Pandas（Python原生函数），必须明确指定返回类型，否则TensorFlow无法推断数据集的结构。
• flat_map更高效，因为它直接生成扁平数据集，避免了先创建大张量再拆分的步骤。

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