Source code for trixi.logger.visdom.pytorchvisdomlogger

import atexit
import tempfile
import warnings
from multiprocessing import Process

import matplotlib.pyplot as plt
from matplotlib import cm
import numpy as np
import torch

from torch.autograd import Variable
from torchvision.utils import make_grid

from trixi.util.util import np_make_grid, get_tensor_embedding
from trixi.logger.experiment.pytorchexperimentlogger import PytorchExperimentLogger
from trixi.logger.visdom.numpyvisdomlogger import NumpyVisdomLogger, add_to_queue
from trixi.logger.abstractlogger import convert_params
from trixi.util.pytorchutils import get_guided_image_gradient, get_smooth_image_gradient, get_vanilla_image_gradient


from functools import wraps


[docs]def move_to_cpu(fn): """Decorator to call the process_params method of the class.""" def __process_params(*args, **kwargs): ### convert args args = (a.detach().cpu() if torch.is_tensor(a) else a for a in args) ### convert kwargs for key, data in kwargs.items(): if torch.is_tensor(data): kwargs[key] = data.detach().cpu() return fn(*args, **kwargs) # # @wraps(f) # def wrapper(*args, **kwargs): # return __process_params(fn, *args, **kwargs) return __process_params
[docs]class PytorchVisdomLogger(NumpyVisdomLogger): """ Visual logger, inherits the NumpyVisdomLogger and plots/ logs pytorch tensors and variables on a Visdom server. """ def __init__(self, *args, **kwargs): super(PytorchVisdomLogger, self).__init__(*args, **kwargs)
[docs] def process_params(self, f, *args, **kwargs): """ Inherited "decorator": convert Pytorch variables and Tensors to numpy arrays. """ ### convert args args = (a.detach().cpu().numpy() if torch.is_tensor(a) else a for a in args) # args = (a.data.cpu().numpy() if isinstance(a, Variable) else a for a in args) ### convert kwargs for key, data in kwargs.items(): # if isinstance(data, Variable): # kwargs[key] = data.detach().cpu().numpy() if torch.is_tensor(data): kwargs[key] = data.detach().cpu().numpy() return f(self, *args, **kwargs)
[docs] @move_to_cpu def plot_model_statistics(self, model, env_appendix="", model_name="", plot_grad=False, **kwargs): """ Plots statstics (mean, std, abs(max)) of the weights or the corresponding gradients of a model as a barplot. Args: model: Model with the weights. env_appendix: Visdom environment name appendix model_name: Name of the model (is used as window name). plot_grad: If false plots weight statistics, if true plot the gradients of the weights. """ means = [] stds = [] maxmin = [] legendary = [] for i, (m_param_name, m_param) in enumerate(model.named_parameters()): win_name = "%s_params" % model_name if plot_grad: m_param = m_param.grad win_name = "%s_grad" % model_name if m_param is not None: param_mean = m_param.detach().mean().item() param_std = m_param.detach().std().item() if np.isnan(param_std): param_std = 0 means.append(param_mean) stds.append(param_std) maxmin.append(torch.max(torch.abs(m_param)).item()) legendary.append("%s-%s" % (model_name, m_param_name)) self.show_barplot(name=win_name, array=np.asarray([means, stds, maxmin]), legend=legendary, rownames=["mean", "std", "max"], env_appendix=env_appendix)
[docs] def plot_model_statistics_weights(self, model, env_appendix="", model_name="", **kwargs): """ Plots statstics (mean, std, abs(max)) of the weights of a model as a barplot (uses plot model statistics with plot_grad=False). Args: model: Model with the weights. env_appendix: Visdom environment name appendix model_name: Name of the model (is used as window name). """ self.plot_model_statistics(model=model, env_appendix=env_appendix, model_name=model_name, plot_grad=False)
[docs] def plot_model_statistics_grads(self, model, env_appendix="", model_name="", **kwargs): """ Plots statstics (mean, std, abs(max)) of the gradients of a model as a barplot (uses plot model statistics with plot_grad=True). Args: model: Model with the weights and the corresponding gradients (have to calculated previously). env_appendix: Visdom environment name appendix model_name: Name of the model (is used as window name). """ self.plot_model_statistics(model=model, env_appendix=env_appendix, model_name=model_name, plot_grad=True)
[docs] def plot_model_gradient_flow(self, model, name="model", title=None): """ Plots statstics (mean, std, abs(max)) of the weights or the corresponding gradients of a model as a barplot. Args: model: Model with the weights. env_appendix: Visdom environment name appendix, if none is given, it uses "-histogram". model_name: Name of the model (is used as window name). plot_grad: If false plots weight statistics, if true plot the gradients of the weights. """ ave_grads = [] layers = [] named_parameters = model.named_parameters() for n, p in named_parameters: if (p.requires_grad) and ("bias" not in n): layers.append(n) ave_grads.append(p.grad.abs().mean()) plt.figure() plt.plot(ave_grads, alpha=0.3, color="b") plt.hlines(0, 0, len(ave_grads) + 1, linewidth=1, color="k") plt.xticks(range(0, len(ave_grads), 1), layers, rotation="vertical") plt.xlim(xmin=0, xmax=len(ave_grads)) plt.xlabel("Layers") plt.ylabel("average gradient {}".format(name)) plt.title("Gradient flow") plt.grid(True) self.show_matplot_plt(plt.gcf(), name=name, title=title)
[docs] def plot_mutliple_models_statistics_weights(self, model_dict, env_appendix=None, **kwargs): """ Given models in a dict, plots the weight statistics of the models. Args: model_dict: Dict with models, the key is assumed to be the name, while the value is the model. env_appendix: Visdom environment name appendix """ for model_name, model in model_dict.items(): self.plot_model_statistics_weights(model=model, env_appendix=env_appendix, model_name=model_name)
[docs] def plot_mutliple_models_statistics_grads(self, model_dict, env_appendix="", **kwargs): """ Given models in a dict, plots the gradient statistics of the models. Args: model_dict: Dict with models, the key is assumed to be the name, while the value is the model. env_appendix: Visdom environment name appendix """ for model_name, model in model_dict.items(): self.plot_model_statistics_grads(model=model, env_appendix=env_appendix, model_name=model_name)
[docs] def plot_model_structure(self, model, input_size, name="model_structure", use_cuda=True, delete_tmp_on_close=False, forward_kwargs=None, **kwargs): """ Plots the model structure/ model graph of a pytorch module (this only works correctly with pytorch 0.2.0). Args: model: The graph of this model will be plotted. input_size: Input size of the model (with batch dim). name: The name of the window in the visdom env. use_cuda: Perform model dimension calculations on the gpu (cuda). delete_tmp_on_close: Determines if the tmp file will be deleted on close. If set true, can cause problems due to the multi threadded plotting. """ if not hasattr(input_size[0], "__iter__"): input_size = [input_size, ] if not torch.cuda.is_available(): use_cuda = False if forward_kwargs is None: forward_kwargs = {} def make_dot(output_var, state_dict=None): """ Produces Graphviz representation of Pytorch autograd graph. Blue nodes are the Variables that require grad, orange are Tensors saved for backward in torch.autograd.Function. Args: output_var: output Variable state_dict: dict of (name, parameter) to add names to node that require grad """ from graphviz import Digraph if state_dict is not None: # assert isinstance(params.values()[0], Variable) param_map = {id(v): k for k, v in state_dict.items()} node_attr = dict(style='filled', shape='box', align='left', fontsize='12', ranksep='0.1', height='0.2') dot = Digraph(node_attr=node_attr, graph_attr=dict(size="12,12"), format="svg") seen = set() def size_to_str(size): return '(' + (', ').join(['%d' % v for v in size]) + ')' def add_nodes(var): if var not in seen: if torch.is_tensor(var): dot.node(str(id(var)), size_to_str(var.size()), fillcolor='orange') elif hasattr(var, 'variable'): u = var.variable if state_dict is not None and id(u.data) in param_map: node_name = param_map[id(u.data)] else: node_name = "" node_name = '%s\n %s' % (node_name, size_to_str(u.size())) dot.node(str(id(var)), node_name, fillcolor='lightblue') else: node_name = str(type(var).__name__) if node_name.endswith("Backward"): node_name = node_name[:-8] dot.node(str(id(var)), node_name) seen.add(var) if hasattr(var, 'next_functions'): for u in var.next_functions: if u[0] is not None: dot.edge(str(id(u[0])), str(id(var))) add_nodes(u[0]) if hasattr(var, 'saved_tensors'): for t in var.saved_tensors: dot.edge(str(id(t)), str(id(var))) add_nodes(t) add_nodes(output_var.grad_fn) return dot # Create input inpt_vars = [torch.randn(i_s) for i_s in input_size] if use_cuda: if next(model.parameters()).is_cuda: device = next(model.parameters()).device.index else: device = None inpt_vars = [i_v.cuda(device) for i_v in inpt_vars] model = model.cuda(device) # get output output = model(*inpt_vars, **forward_kwargs) # get temp file to store svg in fp = tempfile.NamedTemporaryFile(suffix=".svg", delete=delete_tmp_on_close) g = make_dot(output, model.state_dict()) try: # Create model graph and store it as svg x = g.render(fp.name[:-4], cleanup=True) # Display model graph in visdom self.show_svg(svg=x, name=name) except Exception as e: warnings.warn("Could not render model, make sure the Graphviz executables are on your system.")
[docs] @move_to_cpu @add_to_queue def show_image_grid(self, tensor, name=None, caption=None, env_appendix="", opts=None, image_args=None, **kwargs): """ Calls the save image grid method (for abstract logger combatibility) Args: images: 4d- tensor (N, C, H, W) name: The name of the window caption: Caption of the generated image grid env_appendix: appendix to the environment name, if used the new env is env+env_appendix opts: opts dict for the ploty/ visdom plot, i.e. can set window size, en/disable ticks,... image_args: Arguments for the tensorvision save image method """ if opts is None: opts = {} if image_args is None: image_args = {} if isinstance(tensor, Variable): tensor = tensor.detach() if torch.is_tensor(tensor): assert torch.is_tensor(tensor), "tensor has to be a pytorch tensor or variable" assert tensor.dim() == 4, "tensor has to have 4 dimensions" if not (tensor.size(1) == 1 or tensor.size(1) == 3): warnings.warn("The 1. dimension (channel) has to be either 1 (gray) or 3 (rgb), taking the first " "dimension now !!!") tensor = tensor[:, 0:1, ] tensor_np = tensor.numpy() grid = np_make_grid(tensor_np, **image_args) image = np.clip(grid * 255, a_min=0, a_max=255) image = image.astype(np.uint8) # grid = make_grid(tensor, **image_args) # image = grid.mul(255).clamp(0, 255).byte().numpy() elif isinstance(tensor, np.ndarray): grid = np_make_grid(tensor, **image_args) image = np.clip(grid * 255, a_min=0, a_max=255) image = image.astype(np.uint8) else: raise ValueError("Tensor has to be a torch tensor or a numpy array") opts = opts.copy() opts.update(dict( title=name, caption=caption )) win = self.vis.image( img=image, win=name, env=self.name + env_appendix, opts=opts ) return win
@convert_params @add_to_queue def show_image_grid_heatmap(self, heatmap, background=None, ratio=0.3, colormap=cm.jet, normalize=True, name="heatmap", caption=None, env_appendix="", opts=None, image_args=None, **kwargs): """ Creates heat map from the given map and if given combines it with the background and then displays results with as image grid. Args: heatmap: 4d- tensor (N, C, H, W), if C = 3, colormap won't be applied. background: 4d- tensor (N, C, H, W) name: The name of the window ratio: The ratio to mix the map with the background (0 = only background, 1 = only map) caption: Caption of the generated image grid env_appendix: appendix to the environment name, if used the new env is env+env_appendix opts: opts dict for the ploty/ visdom plot, i.e. can set window size, en/disable ticks,... image_args: Arguments for the tensorvision save image method """ if opts is None: opts = {} if image_args is None: image_args = {} if "normalize" not in image_args: image_args["normalize"] = normalize # if len(heatmap.shape) != 4: # raise IndexError("'heatmap' must have dimensions BxCxHxW!") map_grid = np_make_grid(heatmap, normalize=normalize) # map_grid.shape is (3, X, Y) if heatmap.shape[1] != 3: map_ = colormap(map_grid[0])[..., :-1].transpose(2, 0, 1) else: # heatmap was already RGB, so don't apply colormap map_ = map_grid if background is not None: img_grid = np_make_grid(background, **image_args) fuse_img = (1.0 - ratio) * img_grid + ratio * map_ else: fuse_img = map_ fuse_img = np.clip(fuse_img * 255, a_min=0, a_max=255).astype(np.uint8) opts = opts.copy() opts.update(dict( title=name, caption=caption )) win = self.vis.image( img=fuse_img, win=name, env=self.name + env_appendix, opts=opts ) return win
[docs] @convert_params def show_embedding(self, tensor, labels=None, name=None, method="tsne", n_dims=2, n_neigh=30, meth_args=None, *args, **kwargs): """ Displays a tensor a an embedding Args: tensor: Tensor to be embedded an then displayed labels: Labels of the entries in the tensor (first dimension) name: The name of the window method: Method used for embedding, options are: tsne, standard, ltsa, hessian, modified, isomap, mds, spectral, umap n_dims: dimensions to embed the data into n_neigh: Neighbour parameter to kind of determin the embedding (see t-SNE for more information) meth_args: Further arguments which can be passed to the embedding method """ if meth_args is None: meth_args = {} def __show_embedding(queue, tensor, labels=None, name=None, method="tsne", n_dims=2, n_neigh=30, **meth_args): emb_data = get_tensor_embedding(tensor, method=method, n_dims=n_dims, n_neigh=n_neigh, **meth_args) vis_task = { "type": "scatterplot", "array": emb_data, "labels": labels, "name": name, "env_appendix": "", "opts": {} } queue.put_nowait(vis_task) p = Process(target=__show_embedding, kwargs=dict(queue=self._queue, tensor=tensor, labels=labels, name=name, method=method, n_dims=n_dims, n_neigh=n_neigh, **meth_args )) atexit.register(p.terminate) p.start()
[docs] @convert_params def show_roc_curve(self, tensor, labels, name, reduce_to_n_samples=None, use_sub_process=False): """ Displays a roc curve given a tensor with scores and the coresponding labels Args: tensor: Tensor with scores (e.g class probability ) labels: Labels of the samples to which the scores match name: The name of the window reduce_to_n_samples: Reduce/ downsample to to n samples for fewer data points use_sub_process: Use a sub process to do the processing """ res_fn = lambda tpr, fpr: self.show_lineplot(tpr, fpr, name=name, opts={"fillarea": True, "webgl": True}) PytorchExperimentLogger.get_roc_curve(tensor=tensor, labels=labels, reduce_to_n_samples=reduce_to_n_samples, use_sub_process=use_sub_process, results_fn=res_fn)
[docs] @convert_params def show_pr_curve(self, tensor, labels, name, reduce_to_n_samples=None, use_sub_process=False): """ Displays a precision recall curve given a tensor with scores and the coresponding labels Args: tensor: Tensor with scores (e.g class probability ) labels: Labels of the samples to which the scores match name: The name of the window reduce_to_n_samples: Reduce/ downsample to to n samples for fewer data points use_sub_process: Use a sub process to do the processing """ res_fn = lambda precision, recall: self.show_lineplot(precision, recall, name=name, opts={"fillarea": True, "webgl": True}) PytorchExperimentLogger.get_pr_curve(tensor=tensor, labels=labels, reduce_to_n_samples=reduce_to_n_samples, use_sub_process=use_sub_process, results_fn=res_fn)
[docs] @convert_params def show_classification_metrics(self, tensor, labels, name, metric=("roc-auc", "pr-score"), use_sub_process=False, tag_name=None): """ Displays some classification metrics as line plots in a graph (similar to show value (also uses show value for the caluclated values)) Args: tensor: Tensor with scores (e.g class probability ) labels: Labels of the samples to which the scores match name: The name of the window metric: List of metrics to calculate. Options are: roc-auc, pr-auc, pr-score, mcc, f1 Returns: """ res_fn = lambda val, name, tag: self.show_value(val, name=name, tag=tag) PytorchExperimentLogger.get_classification_metrics(tensor=tensor, labels=labels, name=name, metric=metric, use_sub_process=use_sub_process, tag_name=tag_name, results_fn=res_fn)
[docs] def show_image_gradient(self, model, inpt, err_fn, grad_type="vanilla", n_runs=20, eps=0.1, abs=False, **image_grid_params): """ Given a model creates calculates the error and backpropagates it to the image and saves it (saliency map). Args: model: The model to be evaluated inpt: Input to the model err_fn: The error function the evaluate the output of the model on grad_type: Gradient calculation method, currently supports (vanilla, vanilla-smooth, guided, guided-smooth) ( the guided backprob can lead to segfaults -.-) n_runs: Number of runs for the smooth variants eps: noise scaling to be applied on the input image (noise is drawn from N(0,1)) abs (bool): Flag, if the gradient should be a absolute value **image_grid_params: Params for make image grid. """ grad = PytorchExperimentLogger.get_input_gradient(model=model, inpt=inpt, err_fn=err_fn, grad_type=grad_type, n_runs=n_runs, eps=eps, abs=abs) self.show_image_grid(grad, **image_grid_params)
[docs] def show_video(self, frame_list=None, name="frames", dim="LxHxWxC", scale=1.0, fps=25): self.vis.video(tensor=np.array(frame_list), dim=dim, opts={'fps': fps})