Merge pull request #35 from greaber/master

allow setting sample rate in add_audio method of SummaryWriter

README.md

@@ -46,12 +46,12 @@ for n_iter in range(100):
                                              "arctanx": np.arctan(n_iter)}, n_iter)
     x = torch.rand(32, 3, 64, 64) # output from network
     if n_iter%10==0:
-        x = vutils.make_grid(x, normalize=True, scale_each=True)   
-        writer.add_image('Image', x, n_iter) 
+        x = vutils.make_grid(x, normalize=True, scale_each=True)
+        writer.add_image('Image', x, n_iter)
         x = torch.zeros(sample_rate*2)
         for i in range(x.size(0)):
             x[i] = np.cos(freqs[n_iter//10]*np.pi*float(i)/float(sample_rate)) # sound amplitude should in [-1, 1]
-        writer.add_audio('myAudio', x, n_iter)
+        writer.add_audio('myAudio', x, n_iter, sample_rate=sample_rate)
         writer.add_text('Text', 'text logged at step:'+str(n_iter), n_iter)
         for name, param in resnet18.named_parameters():
             writer.add_histogram(name, param.clone().cpu().data.numpy(), n_iter)
@@ -70,23 +70,18 @@ writer.close()
 
 `python demo.py`
 
-`tensorboard --logdir runs`  
+`tensorboard --logdir runs`
 
 ## Screenshots
 <img src="screenshots/Demo.gif">
 
 
 ## Tweaks
 To show more images in tensorboard's image tab, just
-modify the hardcoded `event_accumulator` in 
+modify the hardcoded `event_accumulator` in
 `~/anaconda3/lib/python3.6/site-packages/tensorflow/tensorboard/backend/application.py`
 as you wish.
 
-
-## TODO
-audio sample rate
-
-
 ## Reference:
 
 https://github.com/TeamHG-Memex/tensorboard_logger

tensorboardX/writer.py

@@ -250,7 +250,7 @@ def __append_to_scalar_dict(self, tag, scalar_value, global_step,
         if not tag in self.scalar_dict.keys():
             self.scalar_dict[tag] = []
         self.scalar_dict[tag].append([timestamp, global_step, scalar_value])
- 
+
     def add_scalar(self, tag, scalar_value, global_step=None):
         """Add scalar data to summary.
         Args:
@@ -317,18 +317,19 @@ def add_image(self, tag, img_tensor, global_step=None):
             img_tensor: :math:`(3, H, W)`. Use ``torchvision.utils.make_grid()`` to prepare it is a good idea.
         """
         self.file_writer.add_summary(image(tag, img_tensor), global_step)
-    def add_audio(self, tag, snd_tensor, global_step=None):
+    def add_audio(self, tag, snd_tensor, global_step=None, sample_rate=44100):
         """Add audio data to summary.
 
         Args:
             tag (string): Data identifier
             snd_tensor (torch.Tensor): Sound data
             global_step (int): Global step value to record
+            sample_rate (int): sample rate in Hz
 
         Shape:
-            snd_tensor: :math:`(1, L)`. The values should between [-1, 1]. The sample rate is currently fixed at 44100 KHz.
+            snd_tensor: :math:`(1, L)`. The values should between [-1, 1].
         """
-        self.file_writer.add_summary(audio(tag, snd_tensor), global_step)
+        self.file_writer.add_summary(audio(tag, snd_tensor, sample_rate=sample_rate), global_step)
     def add_text(self, tag, text_string, global_step=None):
         """Add text data to summary.
 
@@ -342,7 +343,7 @@ def add_text(self, tag, text_string, global_step=None):
             writer.add_text('lstm', 'This is an lstm', 0)
             writer.add_text('rnn', 'This is an rnn', 10)
 
-        """        
+        """
         self.file_writer.add_summary(text(tag, text_string), global_step)
         if tag not in self.text_tags:
             self.text_tags.append(tag)
@@ -356,20 +357,20 @@ def add_graph(self, model, lastVar):
         # no, let tensorboard handles it and show its warning message.
         """Add graph data to summary.
 
-            To draw the graph, you need a model ``m`` and an input variable ``t`` that have correct size for ``m``. 
-            Say you have runned ``r = m(t)``, then you can use ``writer.add_graph(m, r)`` to save the graph. 
-            By default, the input tensor does not require gradient, therefore it will be omitted when back tracing. 
+            To draw the graph, you need a model ``m`` and an input variable ``t`` that have correct size for ``m``.
+            Say you have runned ``r = m(t)``, then you can use ``writer.add_graph(m, r)`` to save the graph.
+            By default, the input tensor does not require gradient, therefore it will be omitted when back tracing.
             To draw the input node, pass an additional parameter ``requires_grad=True`` when creating the input tensor.
 
         Args:
-            model (torch.nn.Module): model to draw. 
+            model (torch.nn.Module): model to draw.
             lastVar (torch.autograd.Variable): the root node start from.
 
         .. note::
-            This is experimental feature. Graph drawing is based on autograd's backward tracing. 
-            It goes along the ``next_functions`` attribute in a variable recursively, drawing each encountered nodes. 
+            This is experimental feature. Graph drawing is based on autograd's backward tracing.
+            It goes along the ``next_functions`` attribute in a variable recursively, drawing each encountered nodes.
             In some cases, the result is strange. See  https://github.com/lanpa/tensorboard-pytorch/issues/7 and https://github.com/lanpa/tensorboard-pytorch/issues/9
-        """      
+        """
         import torch
         if not hasattr(torch.autograd.Variable, 'grad_fn'):
             print('pytorch version is too old, how about build by yourself?')