DeOldify/fasterai/tensorboard.py

import fastai
from fastai import *
from fastai.vision import *
from fastai.callbacks import *
from fastai.vision.gan import *
from fastai.core import *
import statistics
from .images import ModelImageSet
import torchvision.utils as vutils
from tensorboardX import SummaryWriter


class ModelGraphVisualizer():
    def __init__(self):
        return 
     
    def write_model_graph_to_tensorboard(self, md:DataBunch, model:nn.Module, tbwriter:SummaryWriter):
        try:
            x,y = md.one_batch(DatasetType.Valid, detach=False, denorm=False)
            tbwriter.add_graph(model, x)
        except Exception as e:
            print(("Failed to generate graph for model: {0}. Note that there's an outstanding issue with "
                + "scopes being addressed here:  https://github.com/pytorch/pytorch/pull/12400").format(e))

class ModelHistogramVisualizer():
    def __init__(self):
        return 

    def write_tensorboard_histograms(self, model:nn.Module, iter_count:int, tbwriter:SummaryWriter, name:str='model'):
        try:
            for param_name, param in model.named_parameters():
                tbwriter.add_histogram(name + '/weights/' + param_name, param, iter_count)
        except Exception as e:
            print(("Failed to update histogram for model:  {0}").format(e))


class ModelStatsVisualizer(): 
    def __init__(self):
        return 

    def write_tensorboard_stats(self, model:nn.Module, iter_count:int, tbwriter:SummaryWriter, name:str='model'):
        try:
            gradients = [x.grad  for x in model.parameters() if x.grad is not None]
            gradient_nps = [to_np(x.data) for x in gradients]
    
            if len(gradients) == 0:
                return 

            avg_norm = sum(x.data.norm() for x in gradients)/len(gradients)
            tbwriter.add_scalar(name + '/gradients/avg_norm', avg_norm, iter_count)

            median_norm = statistics.median(x.data.norm() for x in gradients)
            tbwriter.add_scalar(name + '/gradients/median_norm', median_norm, iter_count)

            max_norm = max(x.data.norm() for x in gradients)
            tbwriter.add_scalar(name + '/gradients/max_norm', max_norm, iter_count)

            min_norm = min(x.data.norm() for x in gradients)
            tbwriter.add_scalar(name + '/gradients/min_norm', min_norm, iter_count)

            num_zeros = sum((np.asarray(x)==0.0).sum() for x in  gradient_nps)
            tbwriter.add_scalar(name + '/gradients/num_zeros', num_zeros, iter_count)


            avg_gradient= sum(x.data.mean() for x in gradients)/len(gradients)
            tbwriter.add_scalar(name + '/gradients/avg_gradient', avg_gradient, iter_count)

            median_gradient = statistics.median(x.data.median() for x in gradients)
            tbwriter.add_scalar(name + '/gradients/median_gradient', median_gradient, iter_count)

            max_gradient = max(x.data.max() for x in gradients) 
            tbwriter.add_scalar(name + '/gradients/max_gradient', max_gradient, iter_count)

            min_gradient = min(x.data.min() for x in gradients) 
            tbwriter.add_scalar(name + '/gradients/min_gradient', min_gradient, iter_count)
        except Exception as e:
            print(("Failed to update tensorboard stats for model:  {0}").format(e))

class ImageGenVisualizer():
    def output_image_gen_visuals(self, learn:Learner, trn_batch:Tuple, val_batch:Tuple, iter_count:int, tbwriter:SummaryWriter):
        self._output_visuals(learn=learn, batch=val_batch, iter_count=iter_count, tbwriter=tbwriter, ds_type=DatasetType.Valid)
        self._output_visuals(learn=learn, batch=trn_batch, iter_count=iter_count, tbwriter=tbwriter, ds_type=DatasetType.Train)

    def _output_visuals(self, learn:Learner, batch:Tuple, iter_count:int, tbwriter:SummaryWriter, ds_type: DatasetType):
        image_sets = ModelImageSet.get_list_from_model(learn=learn, batch=batch, ds_type=ds_type)
        self._write_tensorboard_images(image_sets=image_sets, iter_count=iter_count, tbwriter=tbwriter, ds_type=ds_type)
    
    def _write_tensorboard_images(self, image_sets:[ModelImageSet], iter_count:int, tbwriter:SummaryWriter, ds_type: DatasetType):
        try:
            orig_images = []
            gen_images = []
            real_images = []

            for image_set in image_sets:
                orig_images.append(image_set.orig.px)
                gen_images.append(image_set.gen.px)
                real_images.append(image_set.real.px)

            prefix = str(ds_type)

            tbwriter.add_image(prefix + ' orig images', vutils.make_grid(orig_images, normalize=True), iter_count)
            tbwriter.add_image(prefix + ' gen images', vutils.make_grid(gen_images, normalize=True), iter_count)
            tbwriter.add_image(prefix + ' real images', vutils.make_grid(real_images, normalize=True), iter_count)
        except Exception as e:
            print(("Failed to update tensorboard images for model:  {0}").format(e))


#--------Below are what you actually want ot use, in practice----------------#



class LearnerTensorboardWriter(LearnerCallback):
    def __init__(self, learn:Learner, base_dir:Path, name:str, loss_iters:int=25, weight_iters:int=1000, stats_iters:int=1000):
        super().__init__(learn=learn)
        self.base_dir = base_dir
        self.name = name
        log_dir = base_dir/name
        self.tbwriter = SummaryWriter(log_dir=str(log_dir))
        self.loss_iters = loss_iters
        self.weight_iters = weight_iters
        self.stats_iters = stats_iters
        self.weight_vis = ModelHistogramVisualizer()
        self.model_vis = ModelStatsVisualizer() 
        self.data = None

    def _update_batches_if_needed(self):
        #one_batch function is extremely slow.  this is an optimization
        update_batches = self.data is not self.learn.data

        if update_batches:
            self.data = self.learn.data
            self.trn_batch = self.learn.data.one_batch(DatasetType.Train, detach=True, denorm=False, cpu=False)
            self.val_batch = self.learn.data.one_batch(DatasetType.Valid, detach=True, denorm=False, cpu=False)

    def _write_model_stats(self, iteration):
        self.model_vis.write_tensorboard_stats(model=self.learn.model, iter_count=iteration, tbwriter=self.tbwriter) 

    def _write_training_loss(self, iteration, last_loss):
        trn_loss = to_np(last_loss)
        self.tbwriter.add_scalar('/loss/trn_loss', trn_loss, iteration)

    def _write_weight_histograms(self, iteration):
        self.weight_vis.write_tensorboard_histograms(model=self.learn.model, iter_count=iteration, tbwriter=self.tbwriter)

    def _write_val_loss(self, iteration, last_metrics):
        #TODO: Not a fan of this indexing but...what to do?
        val_loss = last_metrics[0]
        self.tbwriter.add_scalar('/loss/val_loss', val_loss, iteration)  
    
    def _write_metrics(self, iteration):
        rec = self.learn.recorder

        for i, name in enumerate(rec.names[3:]):
            if len(rec.metrics) == 0: continue
            if len(rec.metrics[-1:]) == 0: continue
            if len(rec.metrics[-1:][0]) == 0: continue
            value = rec.metrics[-1:][0][i]
            if value is None: continue
            self.tbwriter.add_scalar('/metrics/' + name, to_np(value), iteration) 

    def on_batch_end(self, last_loss, metrics, iteration, **kwargs):
        if iteration==0:
            return

        self._update_batches_if_needed()

        if iteration % self.loss_iters == 0: 
            self._write_training_loss(iteration, last_loss)

        if iteration % self.weight_iters == 0:
            self._write_weight_histograms(iteration)

        if iteration % self.stats_iters == 0:
            self._write_model_stats(iteration)

    def on_epoch_end(self, metrics, last_metrics, iteration, **kwargs):
        self._write_val_loss(iteration, last_metrics)
        self._write_metrics(iteration)


class GANTensorboardWriter(LearnerTensorboardWriter):
    def __init__(self, learn:Learner, base_dir:Path, name:str, loss_iters:int=25, weight_iters:int=1000, 
                stats_iters:int=1000, visual_iters:int=100):
        super().__init__(learn=learn, base_dir=base_dir, name=name, loss_iters=loss_iters, 
                        weight_iters=weight_iters, stats_iters=stats_iters)
        self.visual_iters = visual_iters
        self.img_gen_vis = ImageGenVisualizer()

    #override
    def _write_training_loss(self, iteration, last_loss):
        trainer = self.learn.gan_trainer
        recorder = trainer.recorder

        if len(recorder.losses) > 0:      
            trn_loss = to_np((recorder.losses[-1:])[0])
            self.tbwriter.add_scalar('/loss/trn_loss', trn_loss, iteration)

    #override
    def _write_weight_histograms(self, iteration):
        trainer = self.learn.gan_trainer
        generator = trainer.generator
        critic = trainer.critic

        self.weight_vis.write_tensorboard_histograms(model=generator, iter_count=iteration, tbwriter=self.tbwriter, name='generator')
        self.weight_vis.write_tensorboard_histograms(model=critic, iter_count=iteration, tbwriter=self.tbwriter, name='critic')

    #override
    def _write_model_stats(self, iteration):
        trainer = self.learn.gan_trainer
        generator = trainer.generator
        critic = trainer.critic

        self.model_vis.write_tensorboard_stats(model=generator, iter_count=iteration, tbwriter=self.tbwriter, name='generator')
        self.model_vis.write_tensorboard_stats(model=critic, iter_count=iteration, tbwriter=self.tbwriter, name='critic')

    #override
    def _write_val_loss(self, iteration, last_metrics):
        trainer = self.learn.gan_trainer
        recorder = trainer.recorder 

        if len(recorder.val_losses) > 0:
            val_loss = (recorder.val_losses[-1:])[0]
            self.tbwriter.add_scalar('/loss/val_loss', val_loss, iteration) 


    def _write_images(self, iteration):
        trainer = self.learn.gan_trainer
        recorder = trainer.recorder

        gen_mode = trainer.gen_mode
        trainer.switch(gen_mode=True)
        self.img_gen_vis.output_image_gen_visuals(learn=self.learn, trn_batch=self.trn_batch, val_batch=self.val_batch, 
                                                iter_count=iteration, tbwriter=self.tbwriter)
        trainer.switch(gen_mode=gen_mode)

    def on_batch_end(self, metrics, iteration, **kwargs):
        super().on_batch_end(metrics=metrics, iteration=iteration, **kwargs)

        if iteration==0:
            return

        if iteration % self.visual_iters == 0:
            self._write_images(iteration)

              

class ImageGenTensorboardWriter(LearnerTensorboardWriter):
    def __init__(self, learn:Learner, base_dir:Path, name:str, loss_iters:int=25, weight_iters:int=1000, 
                stats_iters:int=1000, visual_iters:int=100):
        super().__init__(learn=learn, base_dir=base_dir, name=name, loss_iters=loss_iters, weight_iters=weight_iters, 
                        stats_iters=stats_iters)
        self.visual_iters = visual_iters
        self.img_gen_vis = ImageGenVisualizer()

    def _write_images(self, iteration):
        self.img_gen_vis.output_image_gen_visuals(learn=self.learn, trn_batch=self.trn_batch, val_batch=self.val_batch, 
            iter_count=iteration, tbwriter=self.tbwriter)

    def on_batch_end(self, metrics, iteration, **kwargs):
        super().on_batch_end(metrics=metrics, iteration=iteration, **kwargs)

        if iteration==0:
            return

        if iteration % self.visual_iters == 0:
            self._write_images(iteration)