如何将tf.data.Dataset生成的字典元组展平合并为单个字典？

解决方案 #

你可以在tf.data.Dataset.from_tensor_slices调用后链式调用map方法，对每个元组元素做字典合并即可，不需要修改原有创建数据集的逻辑，代码示例如下：

features = {'input_ids': some.tf.tensor, 'input_mask': some.tf.tensor, 'segment_ids': some.tf.tensor}

# 原有创建数据集的代码无需修改
support_dataset = tf.data.Dataset.from_tensor_slices((features, dataset["labels"]))

# 新增一行map操作即可完成元组解包合并
support_dataset = support_dataset.map(lambda feat_dict, label: {
    **feat_dict,
    "labels": label
})

逻辑说明 #

• 原有from_tensor_slices返回的数据集每个元素是(特征字典, 标签张量)格式的元组
• map操作会逐元素处理元组：先解包原有特征字典的所有键值对，再新增labels键对应标签张量，最终返回合并后的单个字典
• 该操作基于TensorFlow图模式执行，不会引入额外性能损耗，完全兼容后续batch、shuffle、prefetch等常用数据集操作。

处理后你可以直接通过element['input_ids']、element['labels']的方式访问对应字段。

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