diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh index 431b805aef0d0..7c9c57ba555e1 100755 --- 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 index 46cc0810a91b0..5067e4aa75b11 100644 --- 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 index 111bba120e420..e817494df9821 100644 --- 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 index dc45df10dc6b5..92c072d30f612 100644 --- 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