[TB] Support writing new style scalar (#53496)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/53496

New style vs old style
b306651ab5/tensorboard/data_compat.py (L49-L53)

Writing in new style can help avoid the cost of migration
b306651ab5/tensorboard/data_compat.py (L46)

----

Test Plan:
buck run caffe2/test:tensorboard

 ---

Reviewed By: edward-io

Differential Revision: D26879076

fbshipit-source-id: 43cfe9e1ca52dad3efc10332715d39f1cc984862

test/expect/TestTensorBoard.test_scalar_new_style.expect

@@ -0,0 +1,12 @@
+value {
+  tag: "test_scalar"
+  tensor {
+    dtype: DT_FLOAT
+    float_val: 1.0
+  }
+  metadata {
+    plugin_data {
+      plugin_name: "scalars"
+    }
+  }
+}

test/test_tensorboard.py

@@ -226,6 +226,7 @@ class TestTensorBoardWriter(BaseTestCase):
             )
             writer.add_scalar('data/scalar_systemtime', 0.1, n_iter)
             writer.add_scalar('data/scalar_customtime', 0.2, n_iter, walltime=n_iter)
+            writer.add_scalar('data/new_style', 0.2, n_iter, new_style=True)
             writer.add_scalars('data/scalar_group', {
                 "xsinx": n_iter * np.sin(n_iter),
                 "xcosx": n_iter * np.cos(n_iter),
@@ -483,6 +484,10 @@ class TestTensorBoardSummary(BaseTestCase):
         mesh = summary.mesh('my_mesh', vertices=v, colors=c, faces=f, config_dict=None)
         self.assertTrue(compare_proto(mesh, self))
 
+    def test_scalar_new_style(self):
+        scalar = summary.scalar('test_scalar', 1.0, new_style=True)
+        self.assertTrue(compare_proto(scalar, self))
+
 def remove_whitespace(string):
     return string.replace(' ', '').replace('\t', '').replace('\n', '')
 

torch/utils/tensorboard/summary.py

@@ -230,7 +230,7 @@ def hparams(hparam_dict=None, metric_dict=None, hparam_domain_discrete=None):
     return exp, ssi, sei
 
 
-def scalar(name, scalar, collections=None):
+def scalar(name, scalar, collections=None, new_style=False):
     """Outputs a `Summary` protocol buffer containing a single scalar value.
     The generated Summary has a Tensor.proto containing the input Tensor.
     Args:
@@ -239,15 +239,30 @@ def scalar(name, scalar, collections=None):
       tensor: A real numeric Tensor containing a single value.
       collections: Optional list of graph collections keys. The new summary op is
         added to these collections. Defaults to `[GraphKeys.SUMMARIES]`.
+      new_style: Whether to use new style (tensor field) or old style (simple_value
+        field). New style could lead to faster data loading.
     Returns:
       A scalar `Tensor` of type `string`. Which contains a `Summary` protobuf.
     Raises:
       ValueError: If tensor has the wrong shape or type.
     """
     scalar = make_np(scalar)
-    assert(scalar.squeeze().ndim == 0), 'scalar should be 0D'
+    assert scalar.squeeze().ndim == 0, "scalar should be 0D"
     scalar = float(scalar)
-    return Summary(value=[Summary.Value(tag=name, simple_value=scalar)])
+    if new_style:
+        plugin_data = SummaryMetadata.PluginData(plugin_name="scalars")
+        smd = SummaryMetadata(plugin_data=plugin_data)
+        return Summary(
+            value=[
+                Summary.Value(
+                    tag=name,
+                    tensor=TensorProto(float_val=[scalar], dtype="DT_FLOAT"),
+                    metadata=smd,
+                )
+            ]
+        )
+    else:
+        return Summary(value=[Summary.Value(tag=name, simple_value=scalar)])
 
 
 def histogram_raw(name, min, max, num, sum, sum_squares, bucket_limits, bucket_counts):

torch/utils/tensorboard/writer.py

@@ -312,7 +312,9 @@ class SummaryWriter(object):
             for k, v in metric_dict.items():
                 w_hp.add_scalar(k, v)
 
-    def add_scalar(self, tag, scalar_value, global_step=None, walltime=None):
+    def add_scalar(
+        self, tag, scalar_value, global_step=None, walltime=None, new_style=False
+    ):
         """Add scalar data to summary.
 
         Args:
@@ -321,7 +323,8 @@ class SummaryWriter(object):
             global_step (int): Global step value to record
             walltime (float): Optional override default walltime (time.time())
               with seconds after epoch of event
-
+            new_style (boolean): Whether to use new style (tensor field) or old
+              style (simple_value field). New style could lead to faster data loading.
         Examples::
 
             from torch.utils.tensorboard import SummaryWriter
@@ -341,8 +344,9 @@ class SummaryWriter(object):
         if self._check_caffe2_blob(scalar_value):
             from caffe2.python import workspace
             scalar_value = workspace.FetchBlob(scalar_value)
-        self._get_file_writer().add_summary(
-            scalar(tag, scalar_value), global_step, walltime)
+
+        summary = scalar(tag, scalar_value, new_style=new_style)
+        self._get_file_writer().add_summary(summary, global_step, walltime)
 
     def add_scalars(self, main_tag, tag_scalar_dict, global_step=None, walltime=None):
         """Adds many scalar data to summary.