Fixed a minor bug in run_experiment and created tests for the EMNIST models file.

PiperOrigin-RevId: 272546207

tensorflow_federated/python/research/baselines/emnist/BUILD

@@ -8,6 +8,14 @@ py_library(
     srcs_version = "PY3",
 )
 
+py_test(
+    name = "models_test",
+    srcs = ["models_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [":models"],
+)
+
 py_binary(
     name = "run_experiment",
     srcs = ["run_experiment.py"],

tensorflow_federated/python/research/baselines/emnist/models.py

@@ -143,6 +143,7 @@ def create_original_fedavg_cnn_model(only_digits=True):
       activation=tf.nn.relu)
 
   model = tf.keras.models.Sequential([
+      tf.keras.layers.Reshape(input_shape=(28 * 28,), target_shape=input_shape),
       conv2d(filters=32, input_shape=input_shape),
       max_pool(),
       conv2d(filters=64),

tensorflow_federated/python/research/baselines/emnist/models_test.py

+# Lint as: python3
+# Copyright 2019, The TensorFlow Federated Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from tensorflow_federated.python.research.baselines.emnist import models
+
+
+class ModelCollectionTest(tf.test.TestCase):
+
+  def test_conv_dropout_only_digits_shape(self):
+    image = tf.random.normal([4, 28 * 28])
+    model = models.create_conv_dropout_model(only_digits=True)
+    logits = model(image)
+    self.assertIsNotNone(logits)
+    self.assertEqual(logits.shape, [4, 10])
+
+  def test_conv_dropout_shape(self):
+    image = tf.random.normal([3, 28 * 28])
+    model = models.create_conv_dropout_model(only_digits=False)
+    logits = model(image)
+
+    self.assertIsNotNone(logits)
+    self.assertEqual(logits.shape, [3, 62])
+
+  def test_2nn_number_of_parameters(self):
+    model = models.create_two_hidden_layer_model(
+        only_digits=True, hidden_units=200)
+
+    # We calculate the number of parameters based on the fact that given densely
+    # connected layers of size n and m with bias units, there are (n+1)m
+    # parameters between these layers. The network above should have layers of
+    # size 28*28, 200, 200, and 10.
+    num_model_params = (28 * 28 + 1) * 200 + 201 * 200 + 201 * 10
+    self.assertEqual(model.count_params(), num_model_params)
+
+  def test_resnet_block_parameters(self):
+    resnet_model_1 = models.create_resnet(only_digits=False, num_blocks=5)
+    resnet_model_2 = models.create_resnet(only_digits=False, num_blocks=9)
+    self.assertLess(resnet_model_1.count_params(),
+                    resnet_model_2.count_params())
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow_federated/python/research/baselines/emnist/run_experiment.py

@@ -176,8 +176,8 @@ class MetricsHook(object):
 
 
 def create_compiled_keras_model():
-  """Create compiled keras model."""
-  model = models.create_keras_model()
+  """Create compiled keras model based on the original FedAvg CNN."""
+  model = models.create_original_fedavg_cnn_model()
 
   model.compile(
       loss=tf.keras.losses.sparse_categorical_crossentropy,