TensorFlow Dataset：获取生成器输出张量形状的正式新版API方法

可直接复用的标准API #

TensorFlow 2.x全正式版本中，无需引入v1兼容逻辑，直接通过数据集实例的内置公开属性element_spec即可完成形状提取。该属性是官方明确设计用来替代旧版output_shapes、output_types的标准接口，和from_generator要求传入的output_signature参数类型完全一致。

实操代码 #

加载tfds数据集的原有逻辑保持不变，直接读取属性即可：

(ds_train, ds_test), ds_info = tfds.load(
    'mnist',
    split=['train', 'test'],
    shuffle_files=True,
    as_supervised=True,
    with_info=True
)

# 非兼容模式获取数据集结构
train_structure = ds_train.element_spec

# 提取可直接用于赋值、形状校验的TensorShape实例
image_tensor_spec, label_tensor_spec = train_structure
image_shape = image_tensor_spec.shape  # 原始未batch的MNIST场景下输出为TensorShape([28, 28, 1])
label_shape = label_tensor_spec.shape  # 对应输出为TensorShape([])

常见问题说明 #

• 为什么直接调用ds_train.output_signature会失败？
  output_signature只是tf.data.Dataset.from_generator的构造入参名称，并非Dataset类对外暴露的公开读取属性，官方从未给Dataset实例设计这个读接口，自然无法正常调用。
• 为什么用tf.data.DatasetSpec(ds_train)拿到的对象无法直接提取形状？
  打印出来的字符串形式只是DatasetSpec的展示输出，实际这个实例的element_spec字段就存储了完整的TensorSpec结构，直接读取该字段就能拿到可直接赋值的TensorShape实例，不需要做任何字符串解析。
• 新旧接口返回值的差异
  tf.compat.v1.data.get_output_shapes的返回结果会自动包含batch()操作生成的动态批次维度（也就是形状中的None占位）；而新版element_spec会完整反映Dataset经过所有链式操作后的输出结构——只要你提前对Dataset做了.batch()处理，element_spec返回的形状会自动带上对应批次维度，和v1接口的结果完全一致，不存在信息缺失。

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