More optimization on tensorboard functionality

fasterai/tensorboard.py

@@ -5,7 +5,7 @@ from fastai.vision import Image
 from fastai.callbacks import LearnerCallback
 from fastai.core import *
 from fastai.torch_core import *
-from threading import Thread
+from threading import Thread, Event
 import time
 from time import sleep
 from queue import Queue
@@ -14,29 +14,57 @@ import torchvision.utils as vutils
 from abc import ABC, abstractmethod
 from tensorboardX import SummaryWriter
 
-class AsyncTBWriter(ABC):
+
+class TBWriteRequest(ABC):
+    def __init__(self, tbwriter: SummaryWriter, iteration:int):
+        super().__init__()
+        self.tbwriter = tbwriter
+        self.iteration = iteration
+
+    @abstractmethod
+    def write(self):
+        pass   
+
+
+# SummaryWriter writes tend to block quite a bit.  This gets around that and greatly boosts performance.
+# Not all tensorboard writes are using this- just the ones that take a long time.  Note that the 
+# SummaryWriter does actually use a threadsafe consumer/producer design ultimately to write to Tensorboard, 
+# so writes done outside of this async loop should be fine.
+class AsyncTBWriter():
     def __init__(self):
         super().__init__()
-        self.exit = False
+        self.stoprequest = Event()
         self.queue = Queue()
-        self.thread = Thread(target=self._queue_processor)
+        self.thread = Thread(target=self._queue_processor, daemon=True)
         self.thread.start()
 
+    def request_write(self, request: TBWriteRequest):
+        if self.stoprequest.isSet():
+            raise Exception('Close was already called!  Cannot perform this operation.')
+        self.queue.put(request)
+
     def _queue_processor(self):
-        while not self.exit:
+        while not self.stoprequest.isSet():
             while not self.queue.empty():
                 request = self.queue.get()
-                self._write_async(request)
-            sleep(0.1)
+                request.write()
+            sleep(0.2)
 
-    @abstractmethod
-    def _write_async(self, request):
+    #Provided this to stop thread explicitly or by context management (with statement) but thread should end on its own 
+    # upon program exit, due to being a daemon.  So using this is probably unecessary.
+    def close(self):
+        self.stoprequest.set()
+        self.thread.join()
+
+    def __enter__(self):
+        # Nothing to do, thread already started.  Could start thread here to enforce use of context manager 
+        # (but that sounds like a pain and a bit unweildy and unecessary for actual usage)
         pass
 
-    def __del__(self):
-        self.exit = True
-        self.thread.join()
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
 
+asyncTBWriter = AsyncTBWriter() 
 
 class ModelImageSet():
     @staticmethod
@@ -58,144 +86,141 @@ class ModelImageSet():
         self.real = real
         self.gen = gen
 
-#TODO:  There aren't any callbacks using this yet.  Not sure if we want this included (not sure if it's useful, honestly)
-class ModelGraphTBWriter():
-    def __init__(self):
-        return
-
-    def write_model_graph_to_tensorboard(self, md:DataBunch, model:nn.Module, tbwriter:SummaryWriter):
-        x,y = md.one_batch(ds_type=DatasetType.Valid, detach=False, denorm=False)
-        tbwriter.add_graph(model=model, input_to_model=x)
 
-class HistogramTBRequest():
+class HistogramTBRequest(TBWriteRequest):
     def __init__(self, model:nn.Module, iteration:int, tbwriter:SummaryWriter, name:str):
-        self.params = [(name, values.clone().detach()) for (name, values) in model.named_parameters()]
-        self.iteration = iteration
-        self.tbwriter = tbwriter
+        super().__init__(tbwriter=tbwriter, iteration=iteration)
+        self.params = [(name, values.clone().detach().cpu()) for (name, values) in model.named_parameters()]
         self.name = name
 
-#If this isn't done async then this is sloooooow
-class HistogramTBWriter(AsyncTBWriter):
-    def __init__(self):
-        super().__init__()
-
     # override
-    def _write_async(self, request:HistogramTBRequest):
+    def write(self):
         try:
-            params = request.params
-            iteration = request.iteration
-            tbwriter = request.tbwriter
-            name = request.name
-
-            for param_name, values in params:
-                tag = name + '/weights/' + param_name
-                tbwriter.add_histogram(tag=tag, values=values, global_step=iteration)
+            for param_name, values in self.params:
+                tag = self.name + '/weights/' + param_name
+                self.tbwriter.add_histogram(tag=tag, values=values, global_step=self.iteration)
         except Exception as e:
             print(("Failed to write model histograms to Tensorboard:  {0}").format(e))
 
-    def write_tensorboard_histograms(self, model:nn.Module, iteration:int, tbwriter:SummaryWriter, name:str='model'):
-        request = HistogramTBRequest(model, iteration, tbwriter, name)
-        self.queue.put(request)
-
-#This is pretty speedy- Don't think we need async writes here
-class ModelStatsTBWriter():
+#If this isn't done async then this is sloooooow
+class HistogramTBWriter():
     def __init__(self):
+        super().__init__()
+
+    def write(self, model:nn.Module, iteration:int, tbwriter:SummaryWriter, name:str='model'):
+        request = HistogramTBRequest(model=model, iteration=iteration, tbwriter=tbwriter, name=name)
+        asyncTBWriter.request_write(request)
+
+class ModelStatsTBRequest(TBWriteRequest):
+    def __init__(self, model:nn.Module, iteration:int, tbwriter:SummaryWriter, name:str):
+        super().__init__(tbwriter=tbwriter, iteration=iteration)
+        self.gradients = [x.grad.clone().detach().cpu() for x in model.parameters() if x.grad is not None]
+        self.name = name
         self.gradients_root = '/gradients/'
 
-    def write_tensorboard_stats(self, model:nn.Module, iteration:int, tbwriter:SummaryWriter, name:str='model_stats'):
-        gradients = [x.grad for x in model.parameters() if x.grad is not None]
-        gradient_nps = [to_np(x.data) for x in gradients]
+    # override
+    def write(self):
+        try:
+            if len(self.gradients) == 0: return
+
+            gradient_nps = [to_np(x.data) for x in self.gradients]
+            avg_norm = sum(x.data.norm() for x in self.gradients)/len(self.gradients)
+            self.tbwriter.add_scalar(
+                tag=self.name + self.gradients_root + 'avg_norm', scalar_value=avg_norm, global_step=self.iteration)
 
-        if len(gradients) == 0: return
+            median_norm = statistics.median(x.data.norm() for x in self.gradients)
+            self.tbwriter.add_scalar(
+                tag=self.name + self.gradients_root + 'median_norm', scalar_value=median_norm, global_step=self.iteration)
 
-        avg_norm = sum(x.data.norm() for x in gradients)/len(gradients)
-        tbwriter.add_scalar(
-            tag=name + self.gradients_root + 'avg_norm', scalar_value=avg_norm, global_step=iteration)
+            max_norm = max(x.data.norm() for x in self.gradients)
+            self.tbwriter.add_scalar(
+                tag=self.name + self.gradients_root + 'max_norm', scalar_value=max_norm, global_step=self.iteration)
 
-        median_norm = statistics.median(x.data.norm() for x in gradients)
-        tbwriter.add_scalar(
-            tag=name + self.gradients_root + 'median_norm', scalar_value=median_norm, global_step=iteration)
+            min_norm = min(x.data.norm() for x in self.gradients)
+            self.tbwriter.add_scalar(
+                tag=self.name + self.gradients_root + 'min_norm', scalar_value=min_norm, global_step=self.iteration)
 
-        max_norm = max(x.data.norm() for x in gradients)
-        tbwriter.add_scalar(
-            tag=name + self.gradients_root + 'max_norm', scalar_value=max_norm, global_step=iteration)
+            num_zeros = sum((np.asarray(x) == 0.0).sum() for x in gradient_nps)
+            self.tbwriter.add_scalar(
+                tag=self.name + self.gradients_root + 'num_zeros', scalar_value=num_zeros, global_step=self.iteration)
 
-        min_norm = min(x.data.norm() for x in gradients)
-        tbwriter.add_scalar(
-            tag=name + self.gradients_root + 'min_norm', scalar_value=min_norm, global_step=iteration)
+            avg_gradient = sum(x.data.mean() for x in self.gradients)/len(self.gradients)
+            self.tbwriter.add_scalar(
+                tag=self.name + self.gradients_root + 'avg_gradient', scalar_value=avg_gradient, global_step=self.iteration)
 
-        num_zeros = sum((np.asarray(x) == 0.0).sum() for x in gradient_nps)
-        tbwriter.add_scalar(
-            tag=name + self.gradients_root + 'num_zeros', scalar_value=num_zeros, global_step=iteration)
+            median_gradient = statistics.median(x.data.median() for x in self.gradients)
+            self.tbwriter.add_scalar(
+                tag=self.name + self.gradients_root + 'median_gradient', scalar_value=median_gradient, global_step=self.iteration)
 
-        avg_gradient = sum(x.data.mean() for x in gradients)/len(gradients)
-        tbwriter.add_scalar(
-            tag=name + self.gradients_root + 'avg_gradient', scalar_value=avg_gradient, global_step=iteration)
+            max_gradient = max(x.data.max() for x in self.gradients)
+            self.tbwriter.add_scalar(
+                tag=self.name + self.gradients_root + 'max_gradient', scalar_value=max_gradient, global_step=self.iteration)
 
-        median_gradient = statistics.median(x.data.median() for x in gradients)
-        tbwriter.add_scalar(
-            tag=name + self.gradients_root + 'median_gradient', scalar_value=median_gradient, global_step=iteration)
+            min_gradient = min(x.data.min() for x in self.gradients)
+            self.tbwriter.add_scalar(
+                tag=self.name + self.gradients_root + 'min_gradient', scalar_value=min_gradient, global_step=self.iteration)
+        except Exception as e:
+            print(("Failed to write model stats to Tensorboard:  {0}").format(e))
 
-        max_gradient = max(x.data.max() for x in gradients)
-        tbwriter.add_scalar(
-            tag=name + self.gradients_root + 'max_gradient', scalar_value=max_gradient, global_step=iteration)
 
-        min_gradient = min(x.data.min() for x in gradients)
-        tbwriter.add_scalar(
-            tag=name + self.gradients_root + 'min_gradient', scalar_value=min_gradient, global_step=iteration)
+class ModelStatsTBWriter():
+    def write(self, model:nn.Module, iteration:int, tbwriter:SummaryWriter, name:str='model_stats'):
+        request = ModelStatsTBRequest(model=model, iteration=iteration, tbwriter=tbwriter, name=name)
+        asyncTBWriter.request_write(request)
 
 
-class ImageTBRequest():
+class ImageTBRequest(TBWriteRequest):
     def __init__(self, learn:Learner, batch:Tuple, iteration:int, tbwriter:SummaryWriter, ds_type:DatasetType):
+        super().__init__(tbwriter=tbwriter, iteration=iteration)
         self.image_sets = ModelImageSet.get_list_from_model(learn=learn, batch=batch, ds_type=ds_type)
-        self.iteration = iteration
-        self.tbwriter = tbwriter
         self.ds_type = ds_type
 
-#If this isn't done async then this is noticeably slower
-class ImageTBWriter(AsyncTBWriter):
-    def __init__(self):
-        super().__init__()
-
     # override
-    def _write_async(self, request:ImageTBRequest):
+    def write(self):
         try:
             orig_images = []
             gen_images = []
             real_images = []
 
-            for image_set in request.image_sets:
+            for image_set in self.image_sets:
                 orig_images.append(image_set.orig.px)
                 gen_images.append(image_set.gen.px)
                 real_images.append(image_set.real.px)
 
-            prefix = request.ds_type.name
-            tbwriter = request.tbwriter
-            iteration = request.iteration
-
-            tbwriter.add_image(
-                tag=prefix + ' orig images', img_tensor=vutils.make_grid(orig_images, normalize=True), global_step=iteration)
-            tbwriter.add_image(
-                tag=prefix + ' gen images', img_tensor=vutils.make_grid(gen_images, normalize=True), global_step=iteration)
-            tbwriter.add_image(
-                tag=prefix + ' real images', img_tensor=vutils.make_grid(real_images, normalize=True), global_step=iteration)
+            prefix = self.ds_type.name
+
+            self.tbwriter.add_image(
+                tag=prefix + ' orig images', img_tensor=vutils.make_grid(orig_images, normalize=True), 
+                global_step=self.iteration)
+            self.tbwriter.add_image(
+                tag=prefix + ' gen images', img_tensor=vutils.make_grid(gen_images, normalize=True), 
+                global_step=self.iteration)
+            self.tbwriter.add_image(
+                tag=prefix + ' real images', img_tensor=vutils.make_grid(real_images, normalize=True), 
+                global_step=self.iteration)
         except Exception as e:
             print(("Failed to write images to Tensorboard:  {0}").format(e))
 
-    def write_images(self, learn:Learner, trn_batch:Tuple, val_batch:Tuple, iteration:int, tbwriter:SummaryWriter):
-        self._write_images_for_dstype(learn=learn, batch=val_batch, iteration=iteration,
+#If this isn't done async then this is noticeably slower
+class ImageTBWriter():
+    def __init__(self):
+        super().__init__()
+
+    def write(self, learn:Learner, trn_batch:Tuple, val_batch:Tuple, iteration:int, tbwriter:SummaryWriter):
+        self._write_for_dstype(learn=learn, batch=val_batch, iteration=iteration,
                              tbwriter=tbwriter, ds_type=DatasetType.Valid)
-        self._write_images_for_dstype(learn=learn, batch=trn_batch, iteration=iteration,
+        self._write_for_dstype(learn=learn, batch=trn_batch, iteration=iteration,
                              tbwriter=tbwriter, ds_type=DatasetType.Train)
 
-    def _write_images_for_dstype(self, learn:Learner, batch:Tuple, iteration:int, tbwriter:SummaryWriter, ds_type:DatasetType):
+    def _write_for_dstype(self, learn:Learner, batch:Tuple, iteration:int, tbwriter:SummaryWriter, ds_type:DatasetType):
         request = ImageTBRequest(learn=learn, batch=batch, iteration=iteration, tbwriter=tbwriter, ds_type=ds_type)
-        self.queue.put(request)
+        asyncTBWriter.request_write(request)
+
 
 
 #--------CALLBACKS----------------#
 class LearnerTensorboardWriter(LearnerCallback):
-    def __init__(self, learn:Learner, base_dir:Path, name:str, loss_iters:int=25, hist_iters:int=1000, stats_iters:int=1000):
+    def __init__(self, learn:Learner, base_dir:Path, name:str, loss_iters:int=25, hist_iters:int=500, stats_iters:int=100):
         super().__init__(learn=learn)
         self.base_dir = base_dir
         self.name = name
@@ -208,9 +233,12 @@ class LearnerTensorboardWriter(LearnerCallback):
         self.stats_writer = ModelStatsTBWriter()
         self.data = None
         self.metrics_root = '/metrics/'
+        self._update_batches_if_needed()
 
     def _update_batches_if_needed(self):
-        # one_batch function is extremely slow.  this is an optimization
+        # one_batch function is extremely slow with large datasets.  This is an optimization.
+        # Note that also we want to always show the same batches so we can see changes 
+        # in tensorboard
         update_batches = self.data is not self.learn.data
 
         if update_batches:
@@ -221,7 +249,7 @@ class LearnerTensorboardWriter(LearnerCallback):
                 ds_type=DatasetType.Valid, detach=True, denorm=False, cpu=False)
 
     def _write_model_stats(self, iteration:int):
-        self.stats_writer.write_tensorboard_stats(
+        self.stats_writer.write(
             model=self.learn.model, iteration=iteration, tbwriter=self.tbwriter)
 
     def _write_training_loss(self, iteration:int, last_loss:Tensor):
@@ -230,7 +258,7 @@ class LearnerTensorboardWriter(LearnerCallback):
         self.tbwriter.add_scalar(tag=tag, scalar_value=scalar_value, global_step=iteration)
 
     def _write_weight_histograms(self, iteration:int):
-        self.hist_writer.write_tensorboard_histograms(
+        self.hist_writer.write(
             model=self.learn.model, iteration=iteration, tbwriter=self.tbwriter)
 
     #TODO:  Relying on a specific hardcoded start_idx here isn't great.  Is there a better solution?
@@ -266,8 +294,8 @@ class LearnerTensorboardWriter(LearnerCallback):
 
 # TODO:  We're overriding almost everything here.  Seems like a good idea to question that ("is a" vs "has a")
 class GANTensorboardWriter(LearnerTensorboardWriter):
-    def __init__(self, learn:Learner, base_dir:Path, name:str, loss_iters:int=25, hist_iters:int=1000,
-                 stats_iters:int=1000, visual_iters:int=100):
+    def __init__(self, learn:Learner, base_dir:Path, name:str, loss_iters:int=25, hist_iters:int=500,
+                 stats_iters:int=100, visual_iters:int=100):
         super().__init__(learn=learn, base_dir=base_dir, name=name, loss_iters=loss_iters,
                          hist_iters=hist_iters, stats_iters=stats_iters)
         self.visual_iters = visual_iters
@@ -280,9 +308,9 @@ class GANTensorboardWriter(LearnerTensorboardWriter):
         trainer = self.learn.gan_trainer
         generator = trainer.generator
         critic = trainer.critic
-        self.hist_writer.write_tensorboard_histograms(
+        self.hist_writer.write(
             model=generator, iteration=iteration, tbwriter=self.tbwriter, name='generator')
-        self.hist_writer.write_tensorboard_histograms(
+        self.hist_writer.write(
             model=critic, iteration=iteration, tbwriter=self.tbwriter, name='critic')
 
     # override
@@ -295,12 +323,12 @@ class GANTensorboardWriter(LearnerTensorboardWriter):
         gen_mode = trainer.gen_mode
 
         if gen_mode and not self.gen_stats_updated:
-            self.stats_writer.write_tensorboard_stats(
+            self.stats_writer.write(
                 model=generator, iteration=iteration, tbwriter=self.tbwriter, name='gen_model_stats')
             self.gen_stats_updated = True
 
         if not gen_mode and not self.crit_stats_updated:
-            self.stats_writer.write_tensorboard_stats(
+            self.stats_writer.write(
                 model=critic, iteration=iteration, tbwriter=self.tbwriter, name='crit_model_stats')
             self.crit_stats_updated = True
 
@@ -314,14 +342,14 @@ class GANTensorboardWriter(LearnerTensorboardWriter):
             tag = self.metrics_root + 'train_loss'
             self.tbwriter.add_scalar(tag=tag, scalar_value=scalar_value, global_step=iteration)
 
-    def _write_images(self, iteration:int):
+    def _write(self, iteration:int):
         trainer = self.learn.gan_trainer
         #TODO:  Switching gen_mode temporarily seems a bit hacky here.  Certainly not a good side-effect.  Is there a better way?
         gen_mode = trainer.gen_mode
 
         try:
             trainer.switch(gen_mode=True)
-            self.img_gen_vis.write_images(learn=self.learn, trn_batch=self.trn_batch, val_batch=self.val_batch,
+            self.img_gen_vis.write(learn=self.learn, trn_batch=self.trn_batch, val_batch=self.val_batch,
                                                     iteration=iteration, tbwriter=self.tbwriter)
         finally:                                      
             trainer.switch(gen_mode=gen_mode)
@@ -331,7 +359,7 @@ class GANTensorboardWriter(LearnerTensorboardWriter):
         super().on_batch_end(iteration=iteration, **kwargs)
         if iteration == 0: return
         if iteration % self.visual_iters == 0:
-            self._write_images(iteration=iteration)
+            self._write(iteration=iteration)
 
     # override
     def on_backward_end(self, iteration:int, **kwargs):
@@ -349,15 +377,15 @@ class GANTensorboardWriter(LearnerTensorboardWriter):
 
 
 class ImageGenTensorboardWriter(LearnerTensorboardWriter):
-    def __init__(self, learn:Learner, base_dir:Path, name:str, loss_iters:int=25, hist_iters:int=1000,
-                 stats_iters: int = 1000, visual_iters: int = 100):
+    def __init__(self, learn:Learner, base_dir:Path, name:str, loss_iters:int=25, hist_iters:int=500,
+                 stats_iters: int = 100, visual_iters: int = 100):
         super().__init__(learn=learn, base_dir=base_dir, name=name, loss_iters=loss_iters, hist_iters=hist_iters,
                          stats_iters=stats_iters)
         self.visual_iters = visual_iters
         self.img_gen_vis = ImageTBWriter()
 
-    def _write_images(self, iteration:int):
-        self.img_gen_vis.write_images(learn=self.learn, trn_batch=self.trn_batch, val_batch=self.val_batch,
+    def _write(self, iteration:int):
+        self.img_gen_vis.write(learn=self.learn, trn_batch=self.trn_batch, val_batch=self.val_batch,
                                                   iteration=iteration, tbwriter=self.tbwriter)
 
     # override
@@ -366,4 +394,4 @@ class ImageGenTensorboardWriter(LearnerTensorboardWriter):
         if iteration == 0: return
 
         if iteration % self.visual_iters == 0:
-            self._write_images(iteration=iteration)
+            self._write(iteration=iteration)