fastai--fastai/fastai/callback/tensorboard.py

"""Integration with [tensorboard](https://www.tensorflow.org/tensorboard)

Docs: https://docs.fast.ai/callback.tensorboard.html.md"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/70a_callback.tensorboard.ipynb.

# %% auto #0
__all__ = ['TensorBoardBaseCallback', 'TensorBoardCallback', 'TensorBoardProjectorCallback', 'projector_word_embeddings',
           'tensorboard_log']

# %% ../../nbs/70a_callback.tensorboard.ipynb #26870397
from ..basics import *

# %% ../../nbs/70a_callback.tensorboard.ipynb #54b18d89
import tensorboard
from torch.utils.tensorboard import SummaryWriter
from .fp16 import ModelToHalf
from .hook import hook_output

# %% ../../nbs/70a_callback.tensorboard.ipynb #b94e6880
class TensorBoardBaseCallback(Callback):
    order = Recorder.order+1
    "Base class for tensorboard callbacks"
    def __init__(self): self.run_projector = False
        
    def after_pred(self):
        if self.run_projector: self.feat = _add_projector_features(self.learn, self.h, self.feat)
    
    def after_validate(self):
        if not self.run_projector: return
        self.run_projector = False
        self._remove()
        _write_projector_embedding(self.learn, self.writer, self.feat)
            
    def after_fit(self): 
        if self.run: self.writer.close()
        
    def _setup_projector(self):
        self.run_projector = True
        self.h = hook_output(self.learn.model[1][1] if not self.layer else self.layer)
        self.feat = {}
        
    def _setup_writer(self): self.writer = SummaryWriter(log_dir=self.log_dir)
    def __del__(self): self._remove()
    def _remove(self):
        if getattr(self, 'h', None): self.h.remove()

# %% ../../nbs/70a_callback.tensorboard.ipynb #261d75bf
class TensorBoardCallback(TensorBoardBaseCallback):
    "Saves model topology, losses & metrics for tensorboard and tensorboard projector during training"
    def __init__(self, log_dir=None, trace_model=True, log_preds=True, n_preds=9, projector=False, layer=None):
        super().__init__()
        store_attr()

    def before_fit(self):
        self.run = not hasattr(self.learn, 'lr_finder') and not hasattr(self, "gather_preds") and rank_distrib()==0
        if not self.run: return
        self._setup_writer()
        if self.trace_model:
            if hasattr(self.learn, 'mixed_precision'):
                raise Exception("Can't trace model in mixed precision, pass `trace_model=False` or don't use FP16.")
            b = self.dls.one_batch()
            self.learn._split(b)
            self.writer.add_graph(self.model, *self.xb)

    def after_batch(self):
        self.writer.add_scalar('train_loss', self.smooth_loss, self.train_iter)
        for i,h in enumerate(self.opt.hypers):
            for k,v in h.items(): self.writer.add_scalar(f'{k}_{i}', v, self.train_iter)

    def after_epoch(self):
        for n,v in zip(self.recorder.metric_names[2:-1], self.recorder.log[2:-1]):
            self.writer.add_scalar(n, v, self.train_iter)
        if self.log_preds:
            b = self.dls.valid.one_batch()
            self.learn.one_batch(0, b)
            preds = getcallable(self.loss_func, 'activation')(self.pred)
            out = getcallable(self.loss_func, 'decodes')(preds)
            x,y,its,outs = self.dls.valid.show_results(b, out, show=False, max_n=self.n_preds)
            tensorboard_log(x, y, its, outs, self.writer, self.train_iter)
            
    def before_validate(self):
        if self.projector: self._setup_projector()

# %% ../../nbs/70a_callback.tensorboard.ipynb #8b9c90fc
class TensorBoardProjectorCallback(TensorBoardBaseCallback):
    "Extracts and exports image featuers for tensorboard projector during inference"
    def __init__(self, log_dir=None, layer=None):
        super().__init__()
        store_attr()
    
    def before_fit(self):
        self.run = not hasattr(self.learn, 'lr_finder') and hasattr(self, "gather_preds") and rank_distrib()==0
        if not self.run: return
        self._setup_writer()

    def before_validate(self):
        self._setup_projector()

# %% ../../nbs/70a_callback.tensorboard.ipynb #a8757c5a
def _write_projector_embedding(learn, writer, feat):
    lbls = [learn.dl.vocab[l] for l in feat['lbl']] if getattr(learn.dl, 'vocab', None) else None     
    vecs = feat['vec'].squeeze()
    writer.add_embedding(vecs, metadata=lbls, label_img=feat['img'], global_step=learn.train_iter)

# %% ../../nbs/70a_callback.tensorboard.ipynb #fd753d32
def _add_projector_features(learn, hook, feat):
    img = _normalize_for_projector(learn.x)
    first_epoch = True if learn.iter == 0 else False
    feat['vec'] = hook.stored if first_epoch else torch.cat((feat['vec'], hook.stored),0)
    feat['img'] = img           if first_epoch else torch.cat((feat['img'], img),0)
    if getattr(learn.dl, 'vocab', None):
        feat['lbl'] = learn.y if first_epoch else torch.cat((feat['lbl'], learn.y),0)
    return feat

# %% ../../nbs/70a_callback.tensorboard.ipynb #787632a0
def _get_embeddings(model, layer):
    layer = model[0].encoder if layer == None else layer
    return layer.weight

# %% ../../nbs/70a_callback.tensorboard.ipynb #965c546d
@dispatch
def _normalize_for_projector(x:TensorImage):
    # normalize tensor to be between 0-1
    img = x.clone()
    sz = img.shape
    img = img.view(x.size(0), -1)
    img -= img.min(1, keepdim=True)[0]
    img /= img.max(1, keepdim=True)[0]
    img = img.view(*sz)
    return img

# %% ../../nbs/70a_callback.tensorboard.ipynb #1f0a3ae1
from ..text.all import LMLearner, TextLearner

# %% ../../nbs/70a_callback.tensorboard.ipynb #80d48601
def projector_word_embeddings(learn=None, layer=None, vocab=None, limit=-1, start=0, log_dir=None):
    "Extracts and exports word embeddings from language models embedding layers"
    if not layer:
        if   isinstance(learn, LMLearner):   layer = learn.model[0].encoder
        elif isinstance(learn, TextLearner): layer = learn.model[0].module.encoder
    emb = layer.weight
    img = torch.full((len(emb),3,8,8), 0.7)
    vocab = learn.dls.vocab[0] if vocab == None else vocab
    vocab = list(map(lambda x: f'{x}_', vocab))
    writer = SummaryWriter(log_dir=log_dir)
    end = start + limit if limit >= 0 else -1
    writer.add_embedding(emb[start:end], metadata=vocab[start:end], label_img=img[start:end])
    writer.close()

# %% ../../nbs/70a_callback.tensorboard.ipynb #c7528e24
from ..vision.data import *

# %% ../../nbs/70a_callback.tensorboard.ipynb #6247bca8
@dispatch
def tensorboard_log(x:TensorImage, y: TensorCategory, samples, outs, writer, step):
    fig,axs = get_grid(len(samples), return_fig=True)
    for i in range(2):
        axs = [b.show(ctx=c) for b,c in zip(samples.itemgot(i),axs)]
    axs = [r.show(ctx=c, color='green' if b==r else 'red')
            for b,r,c in zip(samples.itemgot(1),outs.itemgot(0),axs)]
    writer.add_figure('Sample results', fig, step)

# %% ../../nbs/70a_callback.tensorboard.ipynb #ac46fb3d
from ..vision.core import TensorPoint,TensorBBox

# %% ../../nbs/70a_callback.tensorboard.ipynb #1073635f
@dispatch
def tensorboard_log(x:TensorImage, y: TensorImageBase|TensorPoint|TensorBBox, samples, outs, writer, step):
    fig,axs = get_grid(len(samples), return_fig=True, double=True)
    for i in range(2):
        axs[::2] = [b.show(ctx=c) for b,c in zip(samples.itemgot(i),axs[::2])]
    for x in [samples,outs]:
        axs[1::2] = [b.show(ctx=c) for b,c in zip(x.itemgot(0),axs[1::2])]
    writer.add_figure('Sample results', fig, step)