tensorflow-federated - 当我们在联合训练数据上调用 IterativeProcess.next 时到底发生了什么？

Question

我浏览了联邦学习教程。我想知道当我们在迭代过程中调用它时 .next 函数是如何工作的。
假设我们有一个列表列表的训练数据。外部列表是客户端列表，内部列表是每个客户端的批量数据。然后，我们创建一个迭代过程，例如联邦平均过程，并初始化状态。
当我们在这个训练数据上调用 IterativeProcess.next 时到底发生了什么。它是否在每一轮中随机从这些数据中获取？还是一次只从每个客户端获取一批数据？
假设我有一个 tf.data.Datasets 列表，每个列表都代表一个客户端数据。如何为下一次联邦学习迭代从此列表中的采样添加一些随机性？
我的数据集的长度不一定相同。当其中一个被完全迭代时，这个数据集是否等待所有其他数据集完全迭代它们的数据？

Answer 1

任意 tff.Computation (如 next )将始终运行整个指定的计算。如果您的 tff.templates.IterativeProcess例如，是 tff.learning.build_federated_averaging_process 的结果，其 next函数将代表一轮联邦平均算法。
联合平均算法在每个本地数据集上运行固定数量的 epochs(为了简单起见，假设为 1 个)的训练，并在服务器上以数据加权的方式平均模型更新以完成一轮 - 参见 Algorithm 1 in the original federated averaging paper对于算法的规范。
现在，关于 TFF 如何表示和执行该算法。来自 build_federated_averaging_process 的文档, next函数具有类型签名:

(<S@SERVER, {B*}@CLIENTS> -> <S@SERVER, T@SERVER>)

TFF 的类型系统将数据集表示为 tff.SequenceType (这就是上面 * 的含义)，所以类型签名的参数中的第二个元素表示一组(技术上是多集)数据集，其元素类型为 B ，放置在客户端。
这在您的示例中的含义如下。您有一个 tf.data.Datasets 的列表，每个都代表每个客户端上的本地数据——您可以将列表视为代表联合放置。在这种情况下，TFF 执行整个指定的计算意味着:TFF 将列表中的每个项目都视为本轮要训练的客户端。根据上面链接的算法，您的数据集列表表示集合 S_t。
TFF 将忠实地执行一轮联邦平均算法， Dataset您的列表中的元素代表为这一轮选择的客户。训练将在每个客户端上运行一个 epoch(并行)；如果数据集具有不同数量的数据，那么您是正确的，每个客户端的训练可能会在不同的时间完成。然而，这是单轮联合平均算法的正确语义，而不是类似算法的参数化，如 Reptile ，它为每个客户端运行固定数量的步骤。
如果您希望选择一个客户端子集来运行一轮训练，这应该在调用 TFF 之前在 Python 中完成，例如:

state = iterative_process.initialize()

# ls is list of datasets
sampled_clients = random.sample(ls, N_CLIENTS)

state = iterative_process.next(state, sampled_clients)

通常，您可以将 Python 运行时视为“实验驱动程序”层——例如，任何客户端的选择都应该发生在这一层。见开头 this answer有关这方面的更多详细信息。