4、TFF数据准备：使用tff.simulation.datasets加载数据集、数据预处理与划分

好，咱们进入实战环节的第一个硬核关卡——数据准备。

说实话，在联邦学习里，数据准备比传统机器学习要麻烦不少。为什么？因为数据不是集中在一个地方，而是分散在各个客户端上。你想想看，传统做法是把所有数据拉到一块儿，洗一洗、切一切，完事。但在联邦场景下，你得模拟出「数据天然分布在各个设备上」的感觉。

TFF 贴心地提供了 tff.simulation.datasets 模块，专门干这个事。我个人习惯，每次开始一个新项目，第一件事就是先跑通数据加载流程，心里才有底。

4.1 内置数据集一览

TFF 内置了几个经典数据集，我列个表，方便你快速了解：

嗯，这里要注意：EMNIST 是我用得最多的数据集。它把每个用户的书写样本作为一个客户端，天然就是非独立同分布的——有的人写「a」很潦草，有的人写「b」很工整。这种数据分布，恰恰是联邦学习要解决的核心难题。

4.2 加载数据集：一行代码搞定

加载代码极其简单，但背后的事情不少。来看：

import tensorflow_federated as tff

# 加载 EMNIST 数据集
emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()

# 看一眼数据结构
print(type(emnist_train))
print(emnist_train.client_ids[0:5])  # 前5个客户端ID

输出结果：

<class 'tff.simulation.datasets.ClientData'>
['00000', '00001', '00002', '00003', '00004']

看到没？返回的是一个 ClientData 对象。它不是普通的 TensorFlow Dataset，而是一个「按客户端索引」的数据容器。你可以通过 client_ids 拿到所有客户端的 ID，然后逐个取出数据。

4.3 数据预处理：别踩坑

数据加载只是第一步。真正的坑，在预处理环节。

我曾经在一个项目里，直接拿原始数据丢给模型训练，结果 loss 死活降不下去。排查了半天，发现是数据没有归一化。你想想看，像素值 0-255 直接喂进神经网络，梯度震荡得跟过山车似的，能收敛才怪。

标准预处理流程，我一般这么写：

def preprocess_dataset(dataset, batch_size=20, shuffle_buffer=1000):
    """对单个客户端的数据进行预处理"""
    def element_fn(element):
        # 归一化：像素值从 [0, 255] 缩放到 [0, 1]
        image = tf.cast(element['pixels'], tf.float32) / 255.0
        # 增加通道维度，变成 [28, 28, 1]
        image = tf.expand_dims(image, axis=-1)
        # 标签保持不变
        label = element['label']
        return {'x': image, 'y': label}
    
    return dataset.map(element_fn).shuffle(shuffle_buffer).batch(batch_size).prefetch(1)

这里有几个关键点：

• 归一化必须做：联邦学习里每个客户端的数据分布可能不同，但归一化是统一的，否则模型学不动。
• prefetch(1)：这个我强烈建议加上。它能让数据加载和模型训练并行起来，GPU 不用干等着。
• shuffle 的 buffer 大小：别设太大，1000 就够。联邦学习里每个客户端的数据量本来就少，shuffle 太大反而浪费内存。

4.4 数据划分：模拟真实联邦场景

数据划分是联邦学习数据准备的核心。说白了，你要模拟出「数据分散在用户手里」的感觉。

TFF 提供了 tff.simulation.datasets.ClientData 的 create_tf_dataset_for_client 方法，可以按客户端取数据：

# 取一个客户端的数据看看
sample_client = emnist_train.client_ids[0]
sample_dataset = emnist_train.create_tf_dataset_for_client(sample_client)
sample_dataset = preprocess_dataset(sample_dataset)

for batch in sample_dataset.take(1):
    print(batch['x'].shape)  # (20, 28, 28, 1)
    print(batch['y'].shape)  # (20,)

但实际训练时，你需要从多个客户端取数据。我一般这么写：

def create_federated_datasets(client_data, client_ids, batch_size=20):
    """从指定客户端列表创建联邦数据集"""
    federated_dataset = []
    for client_id in client_ids:
        # 取单个客户端的数据
        client_ds = client_data.create_tf_dataset_for_client(client_id)
        # 预处理
        client_ds = preprocess_dataset(client_ds, batch_size)
        federated_dataset.append(client_ds)
    return federated_dataset

# 选前 100 个客户端
selected_clients = emnist_train.client_ids[:100]
train_federated = create_federated_datasets(emnist_train, selected_clients)

print(f"共 {len(train_federated)} 个客户端的数据")
print(f"第一个客户端的数据集: {train_federated[0]}")

输出：

共 100 个客户端的数据
第一个客户端的数据集: <BatchDataset shapes: {x: (None, 28, 28, 1), y: (None,)}, types: {x: tf.float32, y: tf.int32}>

4.5 实战：完整的数据准备流水线

最后，我把整个流程串起来，给你一个可以直接用的模板：

import tensorflow as tf
import tensorflow_federated as tff

# 1. 加载数据
emnist_train, emnist_test = tff.simulation.datasets.emnist.load_data()

# 2. 定义预处理函数
def preprocess_fn(dataset, batch_size=20):
    def element_fn(element):
        image = tf.cast(element['pixels'], tf.float32) / 255.0
        image = tf.expand_dims(image, axis=-1)
        label = element['label']
        return {'x': image, 'y': label}
    return dataset.map(element_fn).batch(batch_size).prefetch(1)

# 3. 选择客户端（模拟非独立同分布）
# 我习惯选前 50 个客户端做训练，后 10 个做测试
train_client_ids = emnist_train.client_ids[:50]
test_client_ids = emnist_train.client_ids[50:60]

# 4. 创建联邦数据集
train_federated = [
    preprocess_fn(emnist_train.create_tf_dataset_for_client(cid))
    for cid in train_client_ids
]
test_federated = [
    preprocess_fn(emnist_train.create_tf_dataset_for_client(cid))
    for cid in test_client_ids
]

print(f"训练客户端数: {len(train_federated)}")
print(f"测试客户端数: {len(test_federated)}")
print("数据准备完成，可以开始训练了！")

嗯，到这里，数据准备这块就讲完了。说白了，就是三步：加载、预处理、按客户端划分。但每一步都有细节，尤其是预处理那一步，归一化和 batch_size 的设置，直接影响模型能不能收敛。

我个人建议，你拿到一个新数据集后，先花 10 分钟把数据流水线跑通，打印几个样本看看，确认数据格式正确。这一步省了，后面 debug 的时间会翻倍。

下一章，咱们开始搭建第一个联邦学习模型，到时候就用这套数据流水线。