Add two new attributes to `MnistModel` in TFF tutorial.
This CL is Step 1 of migrating TFF learning model to the new way of binding metrics. It adds two new attributes: `report_local_unfinalized_metrics` and `metric_finalizers`. The two attributes are not being used yet.
},
"outputs": [],
"source": [
+ "from typing import Callable, List, OrderedDict\n",
+ "\n",
"class MnistModel(tff.learning.Model):\n",
"\n",
" def __init__(self):\n",
"\n",
" @property\n",
" def federated_output_computation(self):\n",
- " return aggregate_mnist_metrics_across_clients"
+ " return aggregate_mnist_metrics_across_clients\n",
+ "\n",
+ " @tf.function\n",
+ " def report_local_unfinalized_metrics(\n",
+ " self) -\u003e OrderedDict[str, List[tf.Tensor]]:\n",
+ " \"\"\"Creates an `OrderedDict` of metric names to unfinalized values.\"\"\"\n",
+ " return collections.OrderedDict(\n",
+ " num_examples=[self._variables.num_examples],\n",
+ " loss=[self._variables.loss_sum, self._variables.num_examples],\n",
+ " accuracy=[self._variables.accuracy_sum, self._variables.num_examples])\n",
+ "\n",
+ " def metric_finalizers(\n",
+ " self) -\u003e OrderedDict[str, Callable[[List[tf.Tensor]], tf.Tensor]]:\n",
+ " \"\"\"Creates an `OrderedDict` of metric names to finalizers.\"\"\"\n",
+ " return collections.OrderedDict(\n",
+ " num_examples=tf.function(func=lambda x: x[0]),\n",
+ " loss=tf.function(func=lambda x: x[0] / x[1]),\n",
+ " accuracy=tf.function(func=lambda x: x[0] / x[1]))"
]
},
{
