<p>I think you have the correct pointers for writing a custom federated computation, as well as converting a Keras model to a <code>tff.learning.Model</code>. So we'll focus on pulling a TFF type signature from an existing <code>tff.learning.Model</code>.</p>
<p>Once you have your hands on such a model, you should be able to use <a href="https://www.tensorflow.org/federated/api_docs/python/tff/learning/framework/weights_type_from_model" rel="nofollow noreferrer"><code>tff.learning.framework.weights_type_from_model</code></a> to pull out the appropriate TFF type to use for your custom algorithm.</p>
<p>There is an interesting caveat here: how precisely you <em>use</em> a <code>tff.learning.Model</code> in your custom algorithm is pretty much up to you, and this could affect your desired model weights type. This is unlikely to be the case (likely you will simply be assigning values from incoming tensors to the model variables), so I think we should prefer to avoid going deeper into this caveat.</p>
<p>Finally, a few pointers of end-to-end custom algorithm implementations in TFF:</p>
<ul>
<li>One of the simplest complete examples TFF has is <a href="https://github.com/tensorflow/federated/tree/master/tensorflow_federated/python/examples/simple_fedavg" rel="nofollow noreferrer"><code>simple_fedavg</code></a>, which is totally self-contained and contains instructions for running.</li>
<li>The <a href="https://github.com/google-research/federated/tree/master/optimization" rel="nofollow noreferrer">code</a> for a paper on <a href="https://arxiv.org/pdf/2003.00295.pdf" rel="nofollow noreferrer">Adaptive Federated Optimization</a> contains a <a href="https://github.com/google-research/federated/blob/master/optimization/shared/fed_avg_schedule.py" rel="nofollow noreferrer">handwritten implementation</a> of learning rate decay on the clients in TFF.</li>
<li>A <a href="https://github.com/google-research/federated/blob/master/adaptive_lr_decay/adaptive_fed_avg.py" rel="nofollow noreferrer">similar implementation</a> of adaptive learning rate decay (think Keras' functions to decay learning rate on plateaus) is right next door to the code for AFO.</li>
</ul>
TensorFlow Federated: Keras model with custom learning algorithm
<p><a href="https://www.tensorflow.org/federated/tutorials/federated_learning_for_image_classification" rel="nofollow noreferrer">This</a> tutorial describes how to build a TFF computation from keras model.
<a href="https://www.tensorflow.org/federated/tutorials/custom_federated_algorithms_2" rel="nofollow noreferrer">This</a> tutorial describes how to build a custom TFF computation from scratch, possibly with a custom federated learning algorithm.</p>
<p>What I need is a combination of these: I want to build a custom federated learning algorithm, and I want to use an existing keras model. <strong>Q.</strong> How can it be done?</p>
<p>The second tutorial requires <code>MODEL_TYPE</code> which is based on <code>MODEL_SPEC</code>, but I don't know how to get it. I can see some variables in <code>model.trainable_variables</code> (where <code>model = tff.learning.from_keras_model(keras_model, ...</code>), but I doubt it's what I need.</p>
<p>Of course, I can implement the model by hand (as in the second tutorial), but I want to avoid it.</p>