Performance issue about tf.function about federated-learning-project HOT 2 OPEN

But some variables are depending on the outer function. Code may be more complex if changes are made. Is it necessary to make the change or do you have any other ideas? @luke-who

from federated-learning-project.

Hey! Thank you for raising this. As you mentioned there are variables which depend on the outer function in this case. Here's an idea with the original inner & outer functions stay the same:

# Compile the function outside the loop
batch_train_comp = tff.tf_computation(batch_train)

# Run the loop
model = initial_model
losses = []
train_on_batch = tf.function(batch_train_comp.fn, autograph=False)
for _ in range(5):
  model = train_on_batch(model, sample_batch, 0.1)
  losses.append(batch_loss(model, sample_batch))

In this updated version, batch_train_comp is the compiled TensorFlow function, and train_on_batch is the callable TensorFlow function that wraps around batch_train_comp. The autograph=False argument prevents AutoGraph from converting the loop into a TensorFlow op, so that the function can be reused without retracing. Give it a try and let me know if this works!

from federated-learning-project.