tensorflow-federated 0.28调用联邦平均接口报lambda传参TypeError

在TensorFlow Federated 0.28版本中调用build_weighted_fed_avg_with_optimizer_schedule接口时触发运行时错误，错误栈如下：

Traceback (most recent call last):
  File "/home/Desktop/FL/fedopt.py", line 340, in <module>
    iterative_process = build_weighted_fed_avg_with_optimizer_schedule(
  File "/home/anaconda3/envs/fl/lib/python3.9/site-packages/tensorflow_federated/python/learning/algorithms/fed_avg_with_optimizer_schedule.py", line 276, in build_weighted_fed_avg_with_optimizer_schedule
    client_work = build_scheduled_client_work(model_fn, client_learning_rate_fn,
  File "/home/anaconda3/envs/fl/lib/python3.9/site-packages/tensorflow_federated/python/learning/algorithms/fed_avg_with_optimizer_schedule.py", line 98, in build_scheduled_client_work
    whimsy_optimizer = optimizer_fn(1.0)
TypeError: <lambda>() takes 0 positional arguments but 1 was given

对应的调用代码为：

iterative_process = build_weighted_fed_avg_with_optimizer_schedule(
    model_fn,
    client_learning_rate_fn = lambda x: 0.001,
    client_optimizer_fn=lambda: optimizers.Adam(learning_rate= client_lr, beta_1 = 0.9, beta_2 = 0.999,epsilon = 1e-07),
    server_optimizer_fn=lambda: optimizers.SGD(learning_rate= server_lr), 
    use_experimental_simulation_loop=True)

报错的核心原因是*client_optimizer_fn的函数签名不符合接口要求*：

• 普通FedAvg接口的client_optimizer_fn要求是无参函数，调用后返回优化器实例即可
• 但带学习率调度的build_weighted_fed_avg_with_optimizer_schedule接口中，框架内部会先通过client_learning_rate_fn计算当前轮次的客户端学习率，再将这个学习率值作为入参传给client_optimizer_fn生成对应优化器。接口初始化阶段也会传入1.0作为测试值校验优化器生成逻辑，也就是错误栈中whimsy_optimizer = optimizer_fn(1.0)这一行的逻辑。

当前传入的client_optimizer_fn是无参lambda，框架传入参数时就会触发参数数量不匹配的错误。另外现有写法把Adam的学习率硬编码为外部的client_lr，也会导致配置的client_learning_rate_fn调度逻辑完全不生效。

修改client_optimizer_fn的定义，让它接收一个学习率入参，用传入的学习率初始化Adam优化器即可，修复后的调用代码如下：

iterative_process = build_weighted_fed_avg_with_optimizer_schedule(
    model_fn,
    client_learning_rate_fn = lambda round_num: 0.001, # 入参代表当前训练轮次，后续可按轮次调整学习率实现衰减等逻辑
    # 优化器函数改为接收学习率入参，使用传入的lr初始化优化器
    client_optimizer_fn=lambda lr: optimizers.Adam(learning_rate=lr, beta_1 = 0.9, beta_2 = 0.999,epsilon = 1e-07),
    server_optimizer_fn=lambda: optimizers.SGD(learning_rate= server_lr), 
    use_experimental_simulation_loop=True)

如果后续需要做客户端学习率衰减，只需要修改client_learning_rate_fn的逻辑即可，不需要改动优化器定义。

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