Make mesh plugin tests TF 2.x–compatible (#2560)

Summary:
These just needed to drop down to graph mode in a few places. The
metadata test actually already worked in TF 2.x! :-)

Makes progress toward #1705.

Test Plan:
Tests pass in both TF 1.x and TF 2.x, and the mesh plugin no longer uses
`run_v1_only`:

```
$ git grep run_v1_only '*mesh*' | wc -l
0
```

wchargin-branch: mesh-tests-tf2x

Author: wchargin

tensorboard/plugins/mesh/mesh_plugin_test.py

@@ -44,7 +44,6 @@
   import mock  # pylint: disable=g-import-not-at-top,unused-import
 
 
-@tensorboard_test_util.run_v1_only('requires tf.Session')
 class MeshPluginTest(tf.test.TestCase):
   """Tests for mesh plugin server."""
 
@@ -57,73 +56,71 @@ def setUp(self):
     self.log_dir = self.get_temp_dir()
 
     # Create mesh summary.
-    tf.compat.v1.reset_default_graph()
-    sess = tf.compat.v1.Session()
-    point_cloud = test_utils.get_random_mesh(1000)
-    point_cloud_vertices = tf.compat.v1.placeholder(tf.float32,
-                                                    point_cloud.vertices.shape)
-
-    mesh_no_color = test_utils.get_random_mesh(2000, add_faces=True)
-    mesh_no_color_extended = test_utils.get_random_mesh(2500, add_faces=True)
-    mesh_no_color_vertices = tf.compat.v1.placeholder(
-        tf.float32, [1, None, 3])
-    mesh_no_color_faces = tf.compat.v1.placeholder(tf.int32,
-                                                   [1, None, 3])
-
-    mesh_color = test_utils.get_random_mesh(
-        3000, add_faces=True, add_colors=True)
-    mesh_color_vertices = tf.compat.v1.placeholder(tf.float32,
-                                                   mesh_color.vertices.shape)
-    mesh_color_faces = tf.compat.v1.placeholder(tf.int32,
-                                                mesh_color.faces.shape)
-    mesh_color_colors = tf.compat.v1.placeholder(tf.uint8,
-                                                 mesh_color.colors.shape)
-    self.data = [
-      point_cloud, mesh_no_color, mesh_no_color_extended, mesh_color]
-
-    # In case when name is present and display_name is not, we will reuse name
-    # as display_name. Summaries below intended to test both cases.
-    self.names = ["point_cloud", "mesh_no_color", "mesh_color"]
-    summary.op(
-        self.names[0],
-        point_cloud_vertices,
-        description="just point cloud")
-    summary.op(
-        self.names[1],
-        mesh_no_color_vertices,
-        faces=mesh_no_color_faces,
-        display_name="name_to_display_in_ui",
-        description="beautiful mesh in grayscale")
-    summary.op(
-        self.names[2],
-        mesh_color_vertices,
-        faces=mesh_color_faces,
-        colors=mesh_color_colors,
-        description="mesh with random colors")
-
-    merged_summary_op = tf.compat.v1.summary.merge_all()
-    self.runs = ["bar"]
-    self.steps = 20
-    bar_directory = os.path.join(self.log_dir, self.runs[0])
-    with tensorboard_test_util.FileWriterCache.get(bar_directory) as writer:
-      writer.add_graph(sess.graph)
-      for step in range(self.steps):
-        # Alternate between two random meshes with different number of
-        # vertices.
-        no_color = mesh_no_color if step % 2 == 0 else mesh_no_color_extended
-        with patch.object(time, 'time', return_value=step):
-          writer.add_summary(
-              sess.run(
-                  merged_summary_op,
-                  feed_dict={
-                      point_cloud_vertices: point_cloud.vertices,
-                      mesh_no_color_vertices: no_color.vertices,
-                      mesh_no_color_faces: no_color.faces,
-                      mesh_color_vertices: mesh_color.vertices,
-                      mesh_color_faces: mesh_color.faces,
-                      mesh_color_colors: mesh_color.colors,
-                  }),
-              global_step=step)
+    with tf.compat.v1.Graph().as_default():
+      tf_placeholder = tf.compat.v1.placeholder
+      sess = tf.compat.v1.Session()
+      point_cloud = test_utils.get_random_mesh(1000)
+      point_cloud_vertices = tf_placeholder(
+          tf.float32, point_cloud.vertices.shape
+      )
+
+      mesh_no_color = test_utils.get_random_mesh(2000, add_faces=True)
+      mesh_no_color_extended = test_utils.get_random_mesh(2500, add_faces=True)
+      mesh_no_color_vertices = tf_placeholder(tf.float32, [1, None, 3])
+      mesh_no_color_faces = tf_placeholder(tf.int32, [1, None, 3])
+
+      mesh_color = test_utils.get_random_mesh(
+          3000, add_faces=True, add_colors=True)
+      mesh_color_vertices = tf_placeholder(tf.float32, mesh_color.vertices.shape)
+      mesh_color_faces = tf_placeholder(tf.int32, mesh_color.faces.shape)
+      mesh_color_colors = tf_placeholder(tf.uint8, mesh_color.colors.shape)
+
+      self.data = [
+        point_cloud, mesh_no_color, mesh_no_color_extended, mesh_color]
+
+      # In case when name is present and display_name is not, we will reuse name
+      # as display_name. Summaries below intended to test both cases.
+      self.names = ["point_cloud", "mesh_no_color", "mesh_color"]
+      summary.op(
+          self.names[0],
+          point_cloud_vertices,
+          description="just point cloud")
+      summary.op(
+          self.names[1],
+          mesh_no_color_vertices,
+          faces=mesh_no_color_faces,
+          display_name="name_to_display_in_ui",
+          description="beautiful mesh in grayscale")
+      summary.op(
+          self.names[2],
+          mesh_color_vertices,
+          faces=mesh_color_faces,
+          colors=mesh_color_colors,
+          description="mesh with random colors")
+
+      merged_summary_op = tf.compat.v1.summary.merge_all()
+      self.runs = ["bar"]
+      self.steps = 20
+      bar_directory = os.path.join(self.log_dir, self.runs[0])
+      with tensorboard_test_util.FileWriterCache.get(bar_directory) as writer:
+        writer.add_graph(sess.graph)
+        for step in range(self.steps):
+          # Alternate between two random meshes with different number of
+          # vertices.
+          no_color = mesh_no_color if step % 2 == 0 else mesh_no_color_extended
+          with patch.object(time, 'time', return_value=step):
+            writer.add_summary(
+                sess.run(
+                    merged_summary_op,
+                    feed_dict={
+                        point_cloud_vertices: point_cloud.vertices,
+                        mesh_no_color_vertices: no_color.vertices,
+                        mesh_no_color_faces: no_color.faces,
+                        mesh_color_vertices: mesh_color.vertices,
+                        mesh_color_faces: mesh_color.faces,
+                        mesh_color_colors: mesh_color.colors,
+                    }),
+                global_step=step)
 
     # Start a server that will receive requests.
     self.multiplexer = event_multiplexer.EventMultiplexer({

tensorboard/plugins/mesh/metadata_test.py

@@ -26,7 +26,6 @@
 from tensorboard.util import test_util
 
 
-@test_util.run_v1_only('requires tf.Session')
 class MetadataTest(tf.test.TestCase):
 
   def _create_metadata(self, shape=None):

tensorboard/plugins/mesh/summary.py

@@ -64,7 +64,7 @@ def _get_tensor_summary(
       shape,
       description,
       json_config=json_config)
-  tensor_summary = tf.summary.tensor_summary(
+  tensor_summary = tf.compat.v1.summary.tensor_summary(
       metadata.get_instance_name(name, content_type),
       tensor,
       summary_metadata=tensor_metadata,
@@ -135,7 +135,7 @@ def op(name, vertices, faces=None, colors=None, display_name=None,
                           tensor.content_type, components, json_config,
                           collections))
 
-  all_summaries = tf.summary.merge(
+  all_summaries = tf.compat.v1.summary.merge(
       summaries, collections=collections, name=name)
   return all_summaries
 
@@ -196,9 +196,9 @@ def pb(name,
     tag = metadata.get_instance_name(name, tensor.content_type)
     summaries.append((tag, summary_metadata, tensor_proto))
 
-  summary = tf.Summary()
+  summary = tf.compat.v1.Summary()
   for tag, summary_metadata, tensor_proto in summaries:
-    tf_summary_metadata = tf.SummaryMetadata.FromString(
+    tf_summary_metadata = tf.compat.v1.SummaryMetadata.FromString(
         summary_metadata.SerializeToString())
     summary.value.add(
       tag=tag, metadata=tf_summary_metadata, tensor=tensor_proto)

tensorboard/plugins/mesh/summary_test.py

@@ -28,7 +28,6 @@
 from tensorboard.util import test_util
 
 
-@test_util.run_v1_only('requires tf.Session')
 class MeshSummaryTest(tf.test.TestCase):
 
   def pb_via_op(self, summary_op):
@@ -60,35 +59,37 @@ def test_get_tensor_summary(self):
     description = "my mesh is the best of meshes"
     tensor_data = test_utils.get_random_mesh(100)
     components = 14
-    tensor_summary = summary._get_tensor_summary(
-        name, display_name, description, tensor_data.vertices,
-        plugin_data_pb2.MeshPluginData.VERTEX, components, "", None)
-    with self.test_session():
-      proto = self.pb_via_op(tensor_summary)
-      self.assertEqual("%s_VERTEX" % name, proto.value[0].tag)
-      self.assertEqual(metadata.PLUGIN_NAME,
-                       proto.value[0].metadata.plugin_data.plugin_name)
-      self.assertEqual(components, self.get_components(proto.value[0]))
+    with tf.compat.v1.Graph().as_default():
+      tensor_summary = summary._get_tensor_summary(
+          name, display_name, description, tensor_data.vertices,
+          plugin_data_pb2.MeshPluginData.VERTEX, components, "", None)
+      with self.test_session():
+        proto = self.pb_via_op(tensor_summary)
+        self.assertEqual("%s_VERTEX" % name, proto.value[0].tag)
+        self.assertEqual(metadata.PLUGIN_NAME,
+                         proto.value[0].metadata.plugin_data.plugin_name)
+        self.assertEqual(components, self.get_components(proto.value[0]))
 
   def test_op(self):
     """Tests merged summary with different types of data."""
     name = "my_mesh"
     tensor_data = test_utils.get_random_mesh(
         100, add_faces=True, add_colors=True)
     config_dict = {"foo": 1}
-    tensor_summary = summary.op(
-        name,
-        tensor_data.vertices,
-        faces=tensor_data.faces,
-        colors=tensor_data.colors,
-        config_dict=config_dict)
-    with self.test_session():
-      proto = self.pb_via_op(tensor_summary)
-      self.verify_proto(proto, name)
-      plugin_metadata = metadata.parse_plugin_metadata(
-          proto.value[0].metadata.plugin_data.content)
-      self.assertEqual(
-          json.dumps(config_dict, sort_keys=True), plugin_metadata.json_config)
+    with tf.compat.v1.Graph().as_default():
+      tensor_summary = summary.op(
+          name,
+          tensor_data.vertices,
+          faces=tensor_data.faces,
+          colors=tensor_data.colors,
+          config_dict=config_dict)
+      with self.test_session() as sess:
+        proto = self.pb_via_op(tensor_summary)
+        self.verify_proto(proto, name)
+        plugin_metadata = metadata.parse_plugin_metadata(
+            proto.value[0].metadata.plugin_data.content)
+        self.assertEqual(
+            json.dumps(config_dict, sort_keys=True), plugin_metadata.json_config)
 
   def test_pb(self):
     """Tests merged summary protobuf with different types of data."""