Validation accuracy does not improve during training round - tensorflow-federated

Question

I would like to do a transfer learning with federated learning for image classification, So I write this:
def create_keras_model():
baseModel = tf.keras.models.load_model("./target_model_path", compile=False)
headModel = baseModel.output
layer1 = tf.keras.layers.GlobalAveragePooling2D()(headModel)
layer1 = tf.keras.layers.Dense(128, activation="relu")(layer1)
model_output = tf.keras.layers.Dense(3, activation="softmax")(layer1)
model = tf.keras.Model(inputs=baseModel.input, outputs=model_output)
for layer in baseModel.layers:
layer.trainable = False
return model
keras_model = create_keras_model()
keras_model.compile(
loss=tf.keras.losses.CategoricalCrossentropy(),
metrics=[tf.keras.metrics.CategoricalAccuracy()])
def create_tff_model():
keras_model_clone = tf.keras.models.clone_model(keras_model)
return tff.learning.from_keras_model(
keras_model_clone,
input_spec=input_spec,
loss=tf.keras.losses.CategoricalCrossentropy(),
metrics=[tf.keras.metrics.CategoricalAccuracy()])
fed_avg = tff.learning.build_federated_averaging_process(
model_fn=create_tff_model, server_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=1.0),
client_optimizer_fn=lambda: tf.keras.optimizers.SGD(learning_rate=0.001))
NUM_ROUNDS = 200
state = fed_avg.initialize()
state = tff.learning.state_with_new_model_weights(
state,
trainable_weights=[v.numpy() for v in keras_model.trainable_weights],
non_trainable_weights=[
v.numpy() for v in keras_model.non_trainable_weights
])
def keras_evaluate(state, round_num):
keras_model = create_keras_model()
keras_model.compile(
loss=tf.keras.losses.CategoricalCrossentropy(),
metrics=[tf.keras.metrics.CategoricalAccuracy()])
state.model.assign_weights_to(keras_model)
loss, accuracy = keras_model.evaluate(centralized_test_data, steps=3032 // 2, verbose=1)
for round_num in range(NUM_ROUNDS):
print('Round {r}'.format(r=round_num))
keras_evaluate(state, round_num)
state, metrics = fed_avg.next(state, train_data)
The problem is that validation accuracy, even after 200 rounds, does not exceed 0.6 and the loss value does not decrease (around 0.8) as below:
Round 1
1516/1516 [==============================] - 23s 15ms/step - loss: 0.8254 - categorical_accuracy: 0.6727
Round 2
1516/1516 [==============================] - 23s 15ms/step - loss: 0.8322 - categorical_accuracy: 0.6650
Round 3
1516/1516 [==============================] - 23s 15ms/step - loss: 0.8259 - categorical_accuracy: 0.6700
Round 4
1516/1516 [==============================] - 23s 15ms/step - loss: 0.8216 - categorical_accuracy: 0.6732
Round 5
1516/1516 [==============================] - 23s 15ms/step - loss: 0.8285 - categorical_accuracy: 0.6725
Round 6
1516/1516 [==============================] - 23s 15ms/step - loss: 0.8126 - categorical_accuracy: 0.6808
Round 7
1516/1516 [==============================] - 23s 15ms/step - loss: 0.8210 - categorical_accuracy: 0.6734
Round 8
1516/1516 [==============================] - 23s 15ms/step - loss: 0.8243 - categorical_accuracy: 0.6707
Round 9
1516/1516 [==============================] - 23s 15ms/step - loss: 0.8274 - categorical_accuracy: 0.6722
...
Round 100
1516/1516 [==============================] - 23s 15ms/step - loss: 0.8181 - categorical_accuracy: 0.6778
Can anyone please tell me from what the problem can be ?