如何对包含非均匀tf.RaggedTensor的tf.data.Dataset进行填充？

云上老周

2026-5-14

解决tf.data.Dataset.padded_batch()与RaggedTensor的形状不兼容问题

这个报错的核心原因很清晰：你的输入是一维变长的RaggedTensor（每个样本的形状是(None,)，也就是单个可变长度的词编码序列），但你给padded_shapes设置的是(None, MAX_LENGTH)这种二维结构，两者维度不匹配，自然会抛出形状兼容错误。

正确设置padded_shapes的方法

假设你的Dataset中每个元素的结构是({'inputs': ragged_input_seq, 'dec_inputs': ragged_dec_seq}, {'outputs': ragged_output_seq})，其中每个RaggedTensor都是一维的可变长度序列，那么你只需要针对每个一维序列指定目标填充长度即可。

写法一（简洁版）

直接指定填充后的一维长度：

dataset = dataset.padded_batch(
    BATCH_SIZE,
    padded_shapes=(
        {'inputs': MAX_LENGTH, 'dec_inputs': MAX_LENGTH},
        {'outputs': MAX_LENGTH}
    )
)

写法二（明确形状元组）

如果你习惯用元组表示形状，也可以写成：

dataset = dataset.padded_batch(
    BATCH_SIZE,
    padded_shapes=(
        {'inputs': (MAX_LENGTH,), 'dec_inputs': (MAX_LENGTH,)},
        {'outputs': (MAX_LENGTH,)}
    )
)

这里的逻辑是：每个一维的RaggedTensor只需要填充到MAX_LENGTH的长度，不需要额外的维度层级。None在padded_shapes中通常用来保留批量维度（但padded_batch会自动处理批量维度的生成，所以你只需要关注单个样本内部的形状即可）。

无需tf.keras.preprocessing的替代填充方案

如果你不想用tf.keras.preprocessing的填充步骤，还有两种更直接的方式：

方案1：用RaggedTensor.to_tensor()提前转换

可以在Dataset的map操作中直接把RaggedTensor转换为固定形状的张量，再进行批处理：

def convert_ragged_to_fixed(x):
    return (
        {
            'inputs': x['inputs'].to_tensor(shape=(MAX_LENGTH,)),
            'dec_inputs': x['dec_inputs'].to_tensor(shape=(MAX_LENGTH,))
        },
        {'outputs': x['outputs'].to_tensor(shape=(MAX_LENGTH,))}
    )

# 先转换形状，再批处理
dataset = dataset.map(convert_ragged_to_fixed).batch(BATCH_SIZE)

不过这种方法会把所有样本都强制填充到MAX_LENGTH，不管批次内的实际序列长度，相比padded_batch会浪费一些内存。

方案2：让padded_batch按批次最大长度填充（可选）

如果不需要固定到MAX_LENGTH，而是希望每个批次按内部最长序列填充（更节省内存），可以把padded_shapes设为None，让TensorFlow自动推断每个批次的填充长度：

dataset = dataset.padded_batch(
    BATCH_SIZE,
    padded_shapes=(
        {'inputs': None, 'dec_inputs': None},
        {'outputs': None}
    )
)

如果你之后需要统一到MAX_LENGTH，可以在模型输入层设置input_shape=(MAX_LENGTH,)，或者在后续步骤中截断/填充到目标长度。

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