eXNN.InnerNeuralViz.api

  1import math
  2import matplotlib
  3import matplotlib.pyplot as plt
  4import numpy as np
  5import torch
  6from typing import Dict, List, Optional
  7from sklearn.decomposition import PCA
  8import umap
  9from eXNN.InnerNeuralViz.hook import get_hook
 10
 11
 12def _plot(embedding, labels):
 13    fig, ax = plt.subplots()
 14    ax.scatter(x=embedding[:, 0], y=embedding[:, 1], c=labels)
 15    ax.axis('Off')
 16    plt.close()
 17    return fig
 18
 19
 20def ReduceDim(data: torch.Tensor,
 21              mode: str) -> np.ndarray:
 22    """This function reduces data dimensionality to 2 dimensions.
 23
 24    Args:
 25        data (torch.Tensor): input data of shape NxC1x...xCk,
 26            where N is the number of data points,
 27            C1,...,Ck are dimensions of each data point
 28        mode (str): dimensionality reduction mode (`umap` or `pca`)
 29
 30    Raises:
 31        ValueError: returned if unsupported mode is provided
 32
 33    Returns:
 34        np.ndarray: data projected on a 2d space, of shape Nx2
 35    """
 36
 37    data = data.detach().cpu().numpy().reshape((len(data), -1))
 38    if mode == 'pca':
 39        return PCA(n_components=2).fit_transform(data)
 40    elif mode == 'umap':
 41        return umap.UMAP().fit_transform(data)
 42    else:
 43        raise ValueError(f'Unsupported mode: `{mode}`')
 44
 45
 46def VisualizeNetSpace(model: torch.nn.Module,
 47                      mode: str,
 48                      data: torch.Tensor,
 49                      layers: Optional[List[str]] = None,
 50                      labels: Optional[torch.Tensor] = None,
 51                      chunk_size: Optional[int] = None)\
 52        -> Dict[str, matplotlib.figure.Figure]:
 53    """This function visulizes data latent representations on neural network layers.
 54
 55    Args:
 56        model (torch.nn.Module): neural network
 57        mode (str): dimensionality reduction mode (`umap` or `pca`)
 58        data (torch.Tensor): input data of shape NxC1x...xCk,
 59            where N is the number of data points,
 60            C1,...,Ck are dimensions of each data point
 61        layers (Optional[List[str]], optional): list of layers for visualization.
 62            Defaults to None. If None, visualization for all layers is performed
 63        labels (Optional[torch.Tensor], optional): data labels (colors).
 64            Defaults to None. If None, all points are visualized with the same color
 65        chunk_size (Optional[int], optional): batch size for data processing.
 66            Defaults to None. If None, all data is processed in one batch
 67
 68    Returns:
 69        Dict[str, matplotlib.figure.Figure]: dictionary with latent
 70            representations visualization for each layer
 71    """
 72
 73    if layers is None:
 74        layers = [_[0] for _ in model.named_children()]
 75    if labels is not None:
 76        labels = labels.detach().cpu().numpy()
 77    hooks = {layer: get_hook(model, layer) for layer in layers}
 78    if chunk_size is None:
 79        with torch.no_grad():
 80            out = model(data)
 81        visualizations = {'input': _plot(ReduceDim(data, mode), labels)}
 82        for layer in layers:
 83            visualizations[layer] = _plot(ReduceDim(hooks[layer].fwd, mode), labels)
 84        return visualizations
 85    else:
 86        representations = {layer: [] for layer in layers}
 87        for i in range(math.ceil(len(data) / chunk_size)):
 88            with torch.no_grad():
 89                out = model(data[i * chunk_size:(i + 1) * chunk_size])
 90            for layer in layers:
 91                representations[layer].append(hooks[layer].fwd.detach().cpu())
 92        visualizations = {'input': _plot(ReduceDim(data, mode), labels)}
 93        for layer in layers:
 94            layer_reprs = torch.cat(representations[layer], dim=0)
 95            visualizations[layer] = _plot(ReduceDim(layer_reprs, mode), labels)
 96        return visualizations
 97
 98
 99def get_random_input(dims: List[int]) -> torch.Tensor:
100    """This function generates uniformly distributed tensor of given shape.
101
102    Args:
103        dims (List[int]): required data shape
104
105    Returns:
106        torch.Tensor: uniformly distributed tensor of given shape
107    """
108    return torch.rand(size=dims)