from .writer import SummaryWriter
from multiprocessing import Value
import multiprocessing as mp
global _writer
_writer = None
[docs]class GlobalSummaryWriter(object):
"""A class that implements an event writer that supports concurrent logging and global logging across
different modules.
The GlobalSummaryWriter class provides a set of API to write TensorBoard events from different processes.
The writer instance can be accessed from different processes or modules. Also, the instance maintains
the ``global_step`` value itself so that the interleaved requests to write an event will not conflict
each other. This ensures that the resulting event file is TensorBoard compatible.
With GlobalSummaryWriter, you can easily log the metrics of your parallel-trained model.
The GlobalSummaryWriter and also be used like the ``logging`` module of Python.
See how ``getSummaryWriter`` is used below.
"""
[docs] def __init__(self, logdir=None, comment='', purge_step=None, max_queue=10,
flush_secs=120, filename_suffix='', write_to_disk=True, log_dir=None,
coalesce_process=True, **kwargs):
"""
Initialize a GlobalSummaryWriter. The resulting instance will maintain a monotonically
increasing ``global_step`` for the the event to be written. So there is no need to pass
the global_step when calling its member functions such as ``add_scalar()``.
All arguments for the constructor will be passed to the ordinary ``SummaryWriter.__init__()`` directly.
Examples::
import multiprocessing as mp
import numpy as np
import time
from tensorboardX import GlobalSummaryWriter
w = GlobalSummaryWriter()
def train(x):
w.add_scalar('poolmap/1', x*np.random.randn())
time.sleep(0.05*np.random.randint(0, 10))
w.add_scalar('poolmap/2', x*np.random.randn())
with mp.Pool() as pool:
pool.map(train, range(100))
Expected result:
.. image:: _static/img/tensorboard/add_scalar_global.png
:scale: 50 %
"""
self.smw = SummaryWriter(logdir=logdir, comment=comment, purge_step=purge_step, max_queue=max_queue,
flush_secs=flush_secs, filename_suffix=filename_suffix, write_to_disk=write_to_disk,
log_dir=log_dir)
self.lock = mp.Lock()
self.scalar_tag_to_step = mp.Manager().dict()
self.image_tag_to_step = mp.Manager().dict()
self.histogram_tag_to_step = mp.Manager().dict()
self.text_tag_to_step = mp.Manager().dict()
self.audio_tag_to_step = mp.Manager().dict()
[docs] def add_scalar(self, tag, scalar_value, walltime=None):
"""Add scalar data to summary.
Args:
tag (string): Data identifier
scalar_value (float): Value to save
walltime (float): Optional override default walltime (time.time()) of event
"""
with self.lock:
if tag in self.scalar_tag_to_step:
self.scalar_tag_to_step[tag] += 1
else:
self.scalar_tag_to_step[tag] = 0
self.smw.add_scalar(tag, scalar_value, self.scalar_tag_to_step[tag], walltime)
# def add_histogram(self, tag, values, bins='tensorflow', walltime=None, max_bins=None):
# """Add histogram to summary.
# Args:
# tag (string): Data identifier
# values (torch.Tensor, numpy.array): Values to build histogram
# bins (string): One of {'tensorflow','auto', 'fd', ...}.
# This determines how the bins are made. You can find
# other options in: https://docs.scipy.org/doc/numpy/reference/generated/numpy.histogram.html
# walltime (float): Optional override default walltime (time.time()) of event
# """
# with self.new_tag_mutex.get_lock():
# if tag in self.histogram_tag_to_step:
# self.histogram_tag_to_step[tag] += 1
# else:
# self.histogram_tag_to_step[tag] = 0
# self.smw.add_histogram(tag,
# values,
# self.histogram_tag_to_step[tag],
# bins=bins,
# walltime=walltime,
# max_bins=max_bins)
[docs] def add_image(self, tag, img_tensor, walltime=None, dataformats='CHW'):
"""Add image data to summary.
Note that this requires the ``pillow`` package.
Args:
tag (string): Data identifier
img_tensor (torch.Tensor, numpy.array): An `uint8` or `float`
Tensor of shape `[channel, height, width]` where `channel` is 1, 3, or 4.
The elements in img_tensor can either have values in [0, 1] (float32) or [0, 255] (uint8).
Users are responsible to scale the data in the correct range/type.
walltime (float): Optional override default walltime (time.time()) of event.
dataformats (string): This parameter specifies the meaning of each dimension of the input tensor.
Shape:
img_tensor: Default is :math:`(3, H, W)`. You can use ``torchvision.utils.make_grid()`` to
convert a batch of tensor into 3xHxW format or use ``add_images()`` and let us do the job.
Tensor with :math:`(1, H, W)`, :math:`(H, W)`, :math:`(H, W, 3)` is also suitible as long as
corresponding ``dataformats`` argument is passed. e.g. CHW, HWC, HW.
"""
with self.lock:
if tag in self.image_tag_to_step:
self.image_tag_to_step[tag] += 1
else:
self.image_tag_to_step[tag] = 0
self.smw.add_image(tag, img_tensor, self.image_tag_to_step[tag], walltime=walltime, dataformats=dataformats)
# def add_audio(self, tag, snd_tensor, sample_rate=44100, walltime=None):
# """Add audio data to summary.
# Args:
# tag (string): Data identifier
# snd_tensor (torch.Tensor): Sound data
# sample_rate (int): sample rate in Hz
# walltime (float): Optional override default walltime (time.time()) of event
# Shape:
# snd_tensor: :math:`(1, L)`. The values should lie between [-1, 1].
# """
# with self.new_tag_mutex.get_lock():
# if tag in self.audio_tag_to_step:
# self.audio_tag_to_step[tag] += 1
# else:
# self.audio_tag_to_step[tag] = 0
# self.smw.add_audio(tag, snd_tensor, self.audio_tag_to_step[tag], sample_rate=44100, walltime=walltime)
[docs] def add_text(self, tag, text_string, walltime=None):
"""Add text data to summary.
Args:
tag (string): Data identifier
text_string (string): String to save
walltime (float): Optional override default walltime (time.time()) of event
"""
with self.lock:
if tag in self.text_tag_to_step:
self.text_tag_to_step[tag] += 1
else:
self.text_tag_to_step[tag] = 0
self.smw.add_text(tag, text_string, global_step=self.text_tag_to_step[tag], walltime=walltime)
[docs] @staticmethod
def getSummaryWriter():
"""Get the writer from global namespace.
Examples::
# main.py
import global_1
import global_2
# global1.py
from tensorboardX import GlobalSummaryWriter
writer = GlobalSummaryWriter.getSummaryWriter() # This creates a new instance.
writer.add_text('my_log', 'greeting from global1')
# global2.py
from tensorboardX import GlobalSummaryWriter
writer = GlobalSummaryWriter.getSummaryWriter() # Get the instance in global1.py.
writer.add_text('my_log', 'greeting from global2')
"""
global _writer
if not hasattr(_writer, "smw") or _writer.smw is None:
_writer = GlobalSummaryWriter()
print("Using the global logger in:", _writer.smw.file_writer.get_logdir())
return _writer
@property
def file_writer(self):
return self.smw._get_file_writer()
def close(self):
self.smw.flush()
self.smw.close()