tensorflow-federated 是否支持动态批量大小？答案

【问题标题】：Does tensorflow-federated support dynamic batch size?tensorflow-federated 是否支持动态批量大小？
【发布时间】：2020-04-16 05:09:39
【问题描述】：

tensorflow-federated 是否支持为不同的模拟设备分配不同的 batch-size，并为不同的 epoch 更改 batch-size？

【问题讨论】：

标签： tensorflow tensorflow2.0 tensorflow-federated federated-learning

【解决方案1】：

TFF 确实支持动态批量大小——这是在类型签名级别编码的，在tff.TensorType 的形状属性中。具有None 关联形状的任何维度都将是动态的。如果您有tff.learning.Model，则input_spec 属性应该有None-size 维度，用于您希望动态的任何维度。

确切的正确规范类型签名取决于您希望对这些动态形状做什么。下面是一个简单的例子，可能会说明更多：

假设您有一个 Keras 模型 model 和一个 tff.simulation.ClientData 对象 client_data。 tff.learning.from_keras_model 的 input_spec 参数将直接填充 tff.learning.Model input_spec，因此您希望在此处指定批次维度可以变化：

input_spec = collections.OrderedDict(
    x=tf.TensorSpec(dtype=tf.float32, shape=[None, 784]),
    y=tf.TensorSpec(dtype=tf.int64, shape=[None]),
)

def model_fn():
  tff_model = tff.learning.from_keras_model(
      keras_model=model,
      input_spec=input_spec,
      # other args,...
  )

然后，在你的 Python 驱动训练循环中，你可以在不同的训练轮次中使用不同的批量大小（或者甚至在同一轮训练中），就像这样（假设我们编写了一个名为 _whatever_batch_size_I_want 的函数，它接受轮数作为参数并返回适合该轮的任何批量大小）：

fedavg_process = tff.learning.build_federated_averaging_process(
    model_fn=model_fn, # other args, ...)

state = fedavg_process.initialize()

for k in range(NUM_ROUNDS):
  batch_size = _whatever_batch_size_you_want(k)
  sampled_client_ids = random.choices(
      client_data.client_ids, k=NUM_CLIENTS_PER_ROUND)
  client_datasets = [
      client_data.create_tf_dataset_for_client(x) for x in sampled_client_ids]
  batched_client_datasets = [ds.batch(batch_size) for ds in client_datasets]
  state = fedavg_process.next(state, batched_client_datasets)

如果需要，您可以使用动态形状和输入规范参数来做更有趣的事情；例如，您可以通过将序列维度指定为大小为None 来训练一个采用可变长度序列的序列处理模型。

【讨论】：