Merge pull request #15286 from ddrakard:plot_model_show_activations

PiperOrigin-RevId: 399077719
keras-team · Sep 26, 2021 · 0f3e1ac · 0f3e1ac
2 parents f4dd5a7 + 7658bfc
commit 0f3e1ac
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 deletions.
diff --git a/keras/utils/vis_utils.py b/keras/utils/vis_utils.py
@@ -290,10 +290,8 @@ def format_shape(shape):
             [format_shape(ishape) for ishape in layer.input_shapes])
       else:
         inputlabels = '?'
-      label = '%s\n|{input:|output:}|{{%s}|{%s}}' % (label,
-                                                     inputlabels,
-                                                     outputlabels)
-
+      label = '{%s}|{input:|output:}|{{%s}}|{{%s}}' % (label, inputlabels,
+                                                       outputlabels)
     if not expand_nested or not isinstance(
         layer, functional.Functional):
       node = pydot.Node(layer_id, label=label)

diff --git a/keras/utils/vis_utils_test.py b/keras/utils/vis_utils_test.py
@@ -105,7 +105,20 @@ def test_plot_model_with_add_loss(self):
     except ImportError:
       pass
 
-  def test_plot_model_cnn_with_activations(self):
+  @parameterized.parameters({
+      'show_shapes': False,
+      'show_dtype': False
+  }, {
+      'show_shapes': False,
+      'show_dtype': True
+  }, {
+      'show_shapes': True,
+      'show_dtype': False
+  }, {
+      'show_shapes': True,
+      'show_dtype': True
+  })
+  def test_plot_model_cnn_with_activations(self, show_shapes, show_dtype):
     model = keras.Sequential()
     model.add(
         keras.layers.Conv2D(
@@ -120,8 +133,8 @@ def test_plot_model_cnn_with_activations(self):
       vis_utils.plot_model(
           model,
           to_file=dot_img_file,
-          show_shapes=True,
-          show_dtype=True,
+          show_shapes=show_shapes,
+          show_dtype=show_dtype,
           show_layer_activations=True)
       self.assertTrue(tf.io.gfile.exists(dot_img_file))
       tf.io.gfile.remove(dot_img_file)