TypeError：无法捕获不受支持的类型 tensorflow.python.keras.engine.functional.Functional 的结果

Question

我想在create_keras_model() 内部加载一个预训练网络 所以我写了这个：

def create_keras_model():
  
    baseModel = tf.keras.models.load_model(model_path, compile=False)
 
    headModel = baseModel.output
    model_output = tf.keras.layers.Dense(3, activation="softmax", name="output")(headModel)

    model = tf.keras.Model(inputs=baseModel.input, outputs=model_output)
    
    return model

def model_fn():
    keras_model = create_keras_model()  
    return tff.learning.from_keras_model(keras_model, input_spec = input_spec, loss=tf.keras.losses.CategoricalCrossentropy(), metrics=[tf.keras.metrics.CategoricalAccuracy()])


@tff.tf_computation
def get_weights_from_disk():
   keras_model = create_keras_model()

   return keras_model

@tff.federated_computation
def server_init():
  # There may be state other than weights that needs to get returned from here,
  # as in the implementation of build_federated_averaging_process.
  return tff.federated_eval(get_weights_from_disk(), tff.SERVER)


old_iterproc = tff.learning.build_federated_averaging_process(model_fn=model_fn, server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0), client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.001))
new_iterproc = tff.templates.IterativeProcess(intialize_fn=server_init,
  next_fn=old_iterproc.next)
state = new_iterproc.initialize()

Answer 1

错误来自return 行：

@tff.tf_computation
def get_weights_from_disk():
   keras_model = create_keras_model()
   return keras_model

这是因为 tff.tf_computation 修饰函数只能返回 tf.Tensor 或 tf.Tensor 的结构（list、tuple 或 collections.OrderedDict）。 tf.keras.Model 不是其中之一。

我们可以使用tff.learning.ModelWeights.from_model 仅返回权重结构，而不是返回tf.keras.Model 对象：

@tff.tf_computation
def get_weights_from_disk():
   keras_model = create_keras_model()
   return tff.learning.ModelWeights.from_model(keras_model)

print(get_weights_from_disk.type_signature)

稍后要在 Keras 模型中使用这些权重，可以使用 tff.learing.ModelWeights.assign_to，或者使用 tf.Variable.assign 的 tf.nest.map_structure：

@tff.tf_computation
def train_model(incoming_weights):
  model = create_keras_model()
  destination_weights = tff.learning.ModelWeights.from_model(model)
  tf.nest.map_structure(lambda a, b: a.assign(b),
                        destination_weights, incoming_weights)
  ...