Add review visualization utilities

[boeddeker]

Issue:

I trained a model and want to visualize it in a jupyter notebook. At the moment my workflow is, that I execute the model and call manually some plotting functions, because the review is designed for tensorboard (especially the images are non-obvious, how to print).

Note: Manually calling the plotting functions is better than visualizing the tensorboard visualizing, because tensorboard doesn't know what axis labels and ticks are and how a proper title is formatted, but simply visualizing the review is faster, because the code is already written.

Suggestion:

Add some utilities to visualize entries of the review. e.g.

for k, (data, sample_rate) in review.get('audios', {}).items():
    pb.io.play(data, sample_rate=sample_rate, normalize=False, name=k)

for audios and something like

with pb.visualization.axes_context(columns=4) as axes:
    for k, image in review['images'].items():
        axes.new
        image = np.einsum('chw->hwc', image)[::-1]
        plt.imshow(image, origin='lower')
        plt.title(k)
        plt.grid(False)

for images.

Two proposals for high level functions:

class VisualizeReview:
    def __init__(self, review, trainer=None):
        self.review = review
        if trainer is not None:
            # Ensure, that loss is in review and add loss to scalars
            _, review = trainer._review_to_loss_and_summary(review)
        else:
            review.setdefault('scalars', {})['loss'] = review['loss']
    
    def __call__(self):
        self.scalars()
        self.audios()
        self.images()
    
    def scalars(self):
        display(pd.Series({
            k: pt.utils.to_numpy(v, detach=True)
            for k, v in self.review['scalars'].items()
        }))
    
    def audios(self):
        for k, (data, sample_rate) in self.review.get('audios', {}).items():
            play(data, sample_rate=sample_rate, normalize=False, name=k)
    
    def images(self, columns=4):
        with pb.visualization.axes_context(columns=columns) as axes:
            for k, image in self.review['images'].items():
                axes.new
                image = np.einsum('chw->hwc', image)
                plt.imshow(
                    image,
                    origin='lower',
                )
                plt.title(k)
                plt.grid(False)

VisualizeReview(model_review)()

def visualize_review(
        review,
        trainer=None,
        axes_context_kwargs=dict(columns=4)
):
    from IPython.display import display

    display(pd.Series({
        k: pt.utils.to_numpy(v, detach=True)
        for k, v in review['scalars'].items()
    }))
    
    for k, (data, sample_rate) in review.get('audios', {}).items():
        play(data, sample_rate=sample_rate, normalize=False, name=k)
        
    with pb.visualization.axes_context(**axes_context_kwargs) as axes:
        for k, image in review['images'].items():
            axes.new
            image = np.einsum('chw->hwc', image)
            image = image[::-1]
            plt.imshow(
                image,
                origin='lower',
            )
            plt.title(k)
            plt.grid(False)

visualize_review(model_review)