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

Question

tensorflow-federated 是否支持为不同的模拟设备分配不同的批大小，并为不同的时期改变批大小？

Answer 1

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

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

假设您有一个 Keras 模型model和一个tff.simulation.ClientData对象client_data。的input_spec参数tff.learning.from_keras_model将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。