TFF中如何访问tf.data.Dataset数据集的标签数据

你使用的federated_train_data是TFF标准的联邦数据集结构：列表中每个元素对应单个客户端的本地tf.data.Dataset，数据集内每个样本为(图像张量, 标签张量)的二元组，通过数据集自带的map方法即可拆分出独立的标签数据，具体实现如下： #

核心实现代码 #

1. 拆分指定客户端的图像与标签 #

以你取第一个客户端数据的场景为例，先拆分特征和标签，再转为可直接使用的数组格式：

# 取第一个客户端的本地数据集
client_ds = federated_train_data[0]

# 拆分图像、标签为两个独立数据集
image_ds = client_ds.map(lambda img, label: img)
label_ds = client_ds.map(lambda img, label: label)

# 转为numpy数组，自动对齐样本顺序
sample_count = len(list(client_ds))
images = next(iter(image_ds.batch(sample_count))).numpy()
true_labels = next(iter(label_ds.batch(sample_count))).numpy()

2. 推理并对齐结果 #

不要直接把整个tf.data.Dataset传入predict，否则无法对齐预测结果和真实标签的顺序，用拆分好的图像数组做推理：

# 得到预测概率后取最大概率位为预测标签
pred_result = keras_model.predict(images, verbose=0)
pred_labels = pred_result.argmax(axis=-1)

# 如果做了类别编码，在此处将编码值映射回原始类别名即可
# 示例：label_mapping = {0:"飞机", 1:"汽车", 2:"鸟类"}
# true_label_text = [label_mapping[i] for i in true_labels]
# pred_label_text = [label_mapping[i] for i in pred_labels]

3. 可视化对照 #

遍历图像数组、对应位置的真实标签、预测标签做绘制，即可实现真实类别与预测类别的并排展示效果。 #

注意事项 #

• 如果需要提取所有客户端的标签，遍历federated_train_data列表，对每个客户端的本地数据集重复上述拆分逻辑即可。
• 所有拆分操作必须在数据集完成shuffle、batch、预处理之后执行，否则会出现图像和标签顺序错位的问题。
• 不要尝试直接对联邦数据集做跨客户端的全局索引，TFF的联邦数据集是按客户端独立存储的，必须先定位到单个客户端的数据集再做拆分。

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