TensorFlow Dataset可变第一维度元素的RaggedTensor批量处理异常

我有一个TensorFlow数据集，每个元素是2阶张量，第一维度长度可变、第二维度固定为3（示例形状如(123, 3)、(100, 3)）。以下是模拟该场景的代码：

lengths = tf.random.uniform(shape=(10,), minval=5, maxval=10, dtype=tf.int32)
# 实际代码中此处为文件名数据集
dataset = tf.data.Dataset.from_tensor_slices(lengths)

def fake_read_file(tensor):
    # 实际代码中为文件读取与预处理逻辑，此处生成模拟数据
    dummy_data = tf.convert_to_tensor([[0.2, 0.2, 0.2]])
    return tf.repeat(dummy_data, tensor, axis=0)

# 因文件读取仅支持eager模式，使用tf.py_function
dataset = dataset.map(lambda x: tf.py_function(fake_read_file, inp=[x], Tout=tf.float32))

需求是用RaggedTensor进行批量处理（不填充输入），期望每个批次的形状为(batch_size, None, 3)。

当前采用的方案如下：

dataset = dataset.map(lambda x: tf.RaggedTensor.from_tensor(x))
dataset = dataset.batch(batch_size=2, drop_remainder=True)

但遍历数据集时，得到的批次形状为(batch_size, None, None)，最后一个维度也变成了不规则维度，请问该如何解决？

问题根源在于tf.py_function返回的张量丢失了形状信息——虽然实际第二维度是3，但TensorFlow无法自动推断，导致转成RaggedTensor时最后一维被误判为可变维度。只需在tf.py_function处理后显式设置张量形状，即可解决问题。

修正后的完整代码：

lengths = tf.random.uniform(shape=(10,), minval=5, maxval=10, dtype=tf.int32)
dataset = tf.data.Dataset.from_tensor_slices(lengths)

def fake_read_file(tensor):
    dummy_data = tf.convert_to_tensor([[0.2, 0.2, 0.2]])
    return tf.repeat(dummy_data, tensor, axis=0)

# 关键：给tf.py_function返回的张量设置形状，明确第一维可变、第二维固定为3
dataset = dataset.map(
    lambda x: tf.py_function(
        fake_read_file, 
        inp=[x], 
        Tout=tf.float32
    ).set_shape((None, 3))
)

# 转成RaggedTensor并批量处理
dataset = dataset.map(lambda x: tf.RaggedTensor.from_tensor(x))
dataset = dataset.batch(batch_size=2, drop_remainder=True)

验证结果：

for batch in dataset:
    print(batch.shape)
    # 输出：(2, None, 3)，符合预期

说明 #

set_shape((None, 3))明确告诉TensorFlow：该张量第一维度长度可变，第二维度固定为3。这样转成RaggedTensor时，仅会将第一维作为不规则维度，第二维保持固定，最终批量后的形状就会符合需求。

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