pytorch · Dawars · May 12, 2019 · May 15, 2019 · May 18, 2019 · May 21, 2019
diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh
@@ -45,12 +45,12 @@ if [[ "$BUILD_ENVIRONMENT" != *ppc64le* ]]; then
   # TODO: move this to Docker
   PYTHON_VERSION=$(python -c 'import platform; print(platform.python_version())'|cut -c1)
   echo $PYTHON_VERSION
-  if [[ $PYTHON_VERSION == "2" ]]; then
-    pip install -q https://s3.amazonaws.com/ossci-linux/wheels/tensorboard-1.14.0a0-py2-none-any.whl --user
-  else
-    pip install -q https://s3.amazonaws.com/ossci-linux/wheels/tensorboard-1.14.0a0-py3-none-any.whl --user
-  fi
-
+  # if [[ $PYTHON_VERSION == "2" ]]; then
+  #   pip install -q https://s3.amazonaws.com/ossci-linux/wheels/tensorboard-1.14.0a0-py2-none-any.whl --user
+  # else
+  #   pip install -q https://s3.amazonaws.com/ossci-linux/wheels/tensorboard-1.14.0a0-py3-none-any.whl --user
+  # fi
+  pip install -q tb-nightly --user
   # mypy will fail to install on Python <3.4.  In that case,
   # we just won't run these tests.
   pip install mypy --user || true

diff --git a/docs/source/tensorboard.rst b/docs/source/tensorboard.rst
@@ -96,3 +96,4 @@ Expected result:
    .. automethod:: add_embedding
    .. automethod:: add_pr_curve
    .. automethod:: add_custom_scalars
+   .. automethod:: add_mesh
diff --git a/torch/utils/tensorboard/summary.py b/torch/utils/tensorboard/summary.py
@@ -2,6 +2,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import json
 import logging
 import numpy as np
 import os
@@ -459,3 +460,101 @@ def compute_curve(labels, predictions, num_thresholds=None, weights=None):
     precision = tp / np.maximum(_MINIMUM_COUNT, tp + fp)
     recall = tp / np.maximum(_MINIMUM_COUNT, tp + fn)
     return np.stack((tp, fp, tn, fn, precision, recall))
+
+
+def _get_tensor_summary(name, display_name, description, tensor, content_type, json_config):
+    """Creates a tensor summary with summary metadata.
+
+    Args:
+      name: Uniquely identifiable name of the summary op. Could be replaced by
+        combination of name and type to make it unique even outside of this
+        summary.
+      display_name: Will be used as the display name in TensorBoard.
+        Defaults to `name`.
+      description: A longform readable description of the summary data. Markdown
+        is supported.
+      tensor: Tensor to display in summary.
+      content_type: Type of content inside the Tensor.
+      json_config: A string, JSON-serialized dictionary of ThreeJS classes
+        configuration.
+
+    Returns:
+      Tensor summary with metadata.
+    """
+    import torch
+    from tensorboard.plugins.mesh import metadata
+
+    tensor = torch.as_tensor(tensor)
+
+    tensor_metadata = metadata.create_summary_metadata(
+        name,
+        display_name,
+        content_type,
+        tensor.shape,
+        description,
+        json_config=json_config)
+
+    tensor = TensorProto(dtype='DT_FLOAT',
+                         float_val=tensor.reshape(-1).tolist(),
+                         tensor_shape=TensorShapeProto(dim=[
+                             TensorShapeProto.Dim(size=tensor.shape[0]),
+                             TensorShapeProto.Dim(size=tensor.shape[1]),
+                             TensorShapeProto.Dim(size=tensor.shape[2]),
+                         ]))
+
+    tensor_summary = Summary.Value(
+        tag=metadata.get_instance_name(name, content_type),
+        tensor=tensor,
+        metadata=tensor_metadata,
+    )
+
+    return tensor_summary
+
+
+def _get_json_config(config_dict):
+    """Parses and returns JSON string from python dictionary."""
+    json_config = '{}'
+    if config_dict is not None:
+        json_config = json.dumps(config_dict, sort_keys=True)
+    return json_config
+
+
+# https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/mesh/summary.py
+def mesh(tag, vertices, colors, faces, config_dict, display_name=None, description=None):
+    """Outputs a merged `Summary` protocol buffer with a mesh/point cloud.
+
+      Args:
+        tag: A name for this summary operation.
+        vertices: Tensor of shape `[dim_1, ..., dim_n, 3]` representing the 3D
+          coordinates of vertices.
+        faces: Tensor of shape `[dim_1, ..., dim_n, 3]` containing indices of
+          vertices within each triangle.
+        colors: Tensor of shape `[dim_1, ..., dim_n, 3]` containing colors for each
+          vertex.
+        display_name: If set, will be used as the display name in TensorBoard.
+          Defaults to `name`.
+        description: A longform readable description of the summary data. Markdown
+          is supported.
+        config_dict: Dictionary with ThreeJS classes names and configuration.
+
+      Returns:
+        Merged summary for mesh/point cloud representation.
+      """
+    from tensorboard.plugins.mesh.plugin_data_pb2 import MeshPluginData
+
+    json_config = _get_json_config(config_dict)
+
+    summaries = []
+    tensors = [
+        (vertices, MeshPluginData.VERTEX),
+        (faces, MeshPluginData.FACE),
+        (colors, MeshPluginData.COLOR)
+    ]
+
+    for tensor, content_type in tensors:
+        if tensor is None:
+            continue
+        summaries.append(
+            _get_tensor_summary(tag, display_name, description, tensor, content_type, json_config))
+
+    return Summary(value=summaries)
diff --git a/torch/utils/tensorboard/writer.py b/torch/utils/tensorboard/writer.py
@@ -21,7 +21,7 @@
 from ._utils import figure_to_image
 from .summary import (
     scalar, histogram, histogram_raw, image, audio, text,
-    pr_curve, pr_curve_raw, video, custom_scalars, image_boxes
+    pr_curve, pr_curve_raw, video, custom_scalars, image_boxes, mesh
 )
 
 
@@ -837,6 +837,59 @@ def add_custom_scalars(self, layout):
         """
         self._get_file_writer().add_summary(custom_scalars(layout))
 
+    def add_mesh(self, tag, vertices, colors=None, faces=None, config_dict=None, global_step=None, walltime=None):
+        """Add meshes or 3D point clouds to TensorBoard. The visualization is based on Three.js,
+        so it allows users to interact with the rendered object. Besides the basic definitions
+        such as vertices, faces, users can further provide camera parameter, lighting condition, etc.
+        Please check https://threejs.org/docs/index.html#manual/en/introduction/Creating-a-scene for
+        advanced usage. Note that currently this depends on tb-nightly to show.
+
+        Args:
+            tag (string): Data identifier
+            vertices (torch.Tensor): List of the 3D coordinates of vertices.
+            colors (torch.Tensor): Colors for each vertex
+            faces (torch.Tensor): Indices of vertices within each triangle. (Optional)
+            config_dict: Dictionary with ThreeJS classes names and configuration.
+            global_step (int): Global step value to record
+            walltime (float): Optional override default walltime (time.time())
+              seconds after epoch of event
+
+        Shape:
+            vertices: :math:`(B, N, 3)`. (batch, number_of_vertices, channels)
+
+            colors: :math:`(B, N, 3)`. The values should lie in [0, 255] for type `uint8` or [0, 1] for type `float`.
+
+            faces: :math:`(B, N, 3)`. The values should lie in [0, number_of_vertices] for type `uint8`.
+
+        Examples::
+
+            from torch.utils.tensorboard import SummaryWriter
+            vertices_tensor = torch.as_tensor([
+                [1, 1, 1],
+                [-1, -1, 1],
+                [1, -1, -1],
+                [-1, 1, -1],
+            ], dtype=torch.float).unsqueeze(0)
+            colors_tensor = torch.as_tensor([
+                [255, 0, 0],
+                [0, 255, 0],
+                [0, 0, 255],
+                [255, 0, 255],
+            ], dtype=torch.int).unsqueeze(0)
+            faces_tensor = torch.as_tensor([
+                [0, 2, 3],
+                [0, 3, 1],
+                [0, 1, 2],
+                [1, 3, 2],
+            ], dtype=torch.int).unsqueeze(0)
+
+            writer = SummaryWriter()
+            writer.add_mesh('my_mesh', vertices=vertices_tensor, colors=colors_tensor, faces=faces_tensor)
+
+            writer.close()
+        """
+        self._get_file_writer().add_summary(mesh(tag, vertices, colors, faces, config_dict), global_step, walltime)
+
     def flush(self):
         """Flushes the event file to disk.
         Call this method to make sure that all pending events have been written to