TFF中用set_weights迁移Keras权重时出现AttributeError的解决咨询

问题背景 #

使用TensorFlow Federated构建联邦学习任务时，按以下步骤操作后出现错误：

1. 导入相关库：

import tensorflow as tf
import tensorflow_federated as tff
import collections
import os
import random
import math
import time
import numpy as np
from numpy import sqrt
from numpy.fft import fft, ifft
from numpy.random import rand
import inspect
import tensorflow_probability as tfp
from matplotlib import pyplot as plt
from tensorflow.keras.models import Model
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import  BatchNormalization, AveragePooling2D, MaxPooling2D, Conv2D, Activation, Dropout,Flatten,Input,Dense,concatenate
from tensorflow.keras import layers, initializers
from tensorflow.python.eager import backprop, context, function
from tensorflow.python.framework import constant_op, dtypes, indexed_slices, ops
from tensorflow.python.ops import embedding_ops, math_ops, resource_variable_ops, resources, variables
from tensorflow.python.platform import test
from tensorflow.python.training import gradient_descent

2. 定义Keras模型：

def create_keras_model():
   return tf.keras.models.Sequential([
     tf.keras.layers.Conv2D(filters=64, kernel_size=[5, 5],name='conv2d_1',activation=tf.nn.relu, use_bias=True, bias_initializer =tf.initializers.lecun_normal(seed=137), input_shape=(28 ,28 ,1)),
     tf.keras.layers.MaxPool2D(pool_size=[2,2], strides=2),
     tf.keras.layers.Conv2D(filters=32, kernel_size=[5,5 ],name='conv2d_2',activation=tf.nn.relu, use_bias = True, bias_initializer=tf.initializers.lecun_normal(seed=137)),
     tf.keras.layers.MaxPool2D(pool_size=[2,2], strides=2),
     tf.keras.layers.Reshape(target_shape=(4 * 4 * 32,)),
     tf.keras.layers.Dense(units= 150, activation=tf.nn.relu, use_bias=True, bias_initializer=tf.initializers.lecun_normal(seed=137), name='dense_1'),
     tf.keras.layers.Dense(units=10 , use_bias=True, bias_initializer=tf.initializers.lecun_normal(seed=137), activation=tf.nn.softmax, name='dense_2'   ),
  ])

3. 创建模型实例并编写model_fn传递权重：

net_1 = create_keras_model()

def model_fn():
  global_model = create_keras_model()
  global_model.set_weights(net_1.get_weights())
  return tff.learning.from_keras_model(
      global_model,
      input_spec=preprocessed_example_dataset.element_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

4. 构建加权联邦平均迭代过程时触发错误：

iterative_process = tff.learning.algorithms.build_weighted_fed_avg(
    model_fn, 
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.00))

错误信息 #

AttributeError                            Traceback (most recent call last)
<ipython-input-31-777247538e22> in <module>
      2     model_fn,
      3     client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
----> 4     server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.00))

5 frames
/usr/local/lib/python3.7/dist-packages/keras/engine/training_v1.py in get_weights(self)
    155     """
    156     strategy = (self._distribution_strategy or
--> 157                 self._compile_time_distribution_strategy)
    158     if strategy:
    159       with strategy.scope():

AttributeError: 'Sequential' object has no attribute '_compile_time_distribution_strategy'

解决建议 #

错误根源是：未编译的Keras模型调用get_weights()时缺少_compile_time_distribution_strategy属性，同时TFF要求model_fn必须是无状态的，不能依赖外部作用域的模型对象。

方案一：让模型内部完成权重初始化 #

直接在create_keras_model中固定权重初始化种子，确保每次创建的模型初始权重一致，无需外部传递：

def create_keras_model():
   return tf.keras.models.Sequential([
     tf.keras.layers.Conv2D(filters=64, kernel_size=[5, 5],name='conv2d_1',activation=tf.nn.relu, use_bias=True, bias_initializer =tf.initializers.lecun_normal(seed=137), input_shape=(28 ,28 ,1)),
     tf.keras.layers.MaxPool2D(pool_size=[2,2], strides=2),
     tf.keras.layers.Conv2D(filters=32, kernel_size=[5,5 ],name='conv2d_2',activation=tf.nn.relu, use_bias = True, bias_initializer=tf.initializers.lecun_normal(seed=137)),
     tf.keras.layers.MaxPool2D(pool_size=[2,2], strides=2),
     tf.keras.layers.Reshape(target_shape=(4 * 4 * 32,)),
     tf.keras.layers.Dense(units= 150, activation=tf.nn.relu, use_bias=True, bias_initializer=tf.initializers.lecun_normal(seed=137), name='dense_1'),
     tf.keras.layers.Dense(units=10 , use_bias=True, bias_initializer=tf.initializers.lecun_normal(seed=137), activation=tf.nn.softmax, name='dense_2'   ),
  ])

def model_fn():
  global_model = create_keras_model()
  return tff.learning.from_keras_model(
      global_model,
      input_spec=preprocessed_example_dataset.element_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

方案二：通过TFF初始状态传递外部权重 #

若必须使用外部预训练权重，需先编译模型再获取权重，然后通过TFF的初始状态替换机制传递：

# 创建外部模型并编译（编译后才能正常调用get_weights）
net_1 = create_keras_model()
net_1.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(), metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])
initial_weights = net_1.get_weights()

# 编写无状态的model_fn
def model_fn():
  global_model = create_keras_model()
  return tff.learning.from_keras_model(
      global_model,
      input_spec=preprocessed_example_dataset.element_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

# 构建迭代过程
iterative_process = tff.learning.algorithms.build_weighted_fed_avg(
    model_fn, 
    client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.02),
    server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.00))

# 替换初始状态中的权重
initial_state = iterative_process.initialize()
state_with_custom_weights = initial_state.replace(
    model=initial_state.model._replace(trainable=initial_weights, non_trainable=[])
)
# 后续训练使用state_with_custom_weights作为初始状态

关键注意事项 #

• TFF的model_fn必须是无状态的，不能依赖外部变量，否则会在分布式上下文执行时出错。
• 未编译的Keras模型调用get_weights()会触发该属性错误，因此外部模型需先编译再获取权重。

内容的提问来源于stack exchange，提问作者CA Khan