Function add_ndim_array

slowfast/visualization/tensorboard_vis.py:335–408

Visualize and add tensors of n-dimentionals to a Tensorboard SummaryWriter. Tensors will be visualized as a 2D grid image. Args: writer (SummaryWriter): Tensorboard SummaryWriter. array (tensor): tensor to visualize. name (str): name of the tensor. nrow (

(
    writer,
    array,
    name,
    nrow=None,
    normalize=False,
    global_step=None,
    heat_map=True,
)

def add_ndim_array(
    writer,
    array,
    name,
    nrow=None,
    normalize=False,
    global_step=None,
    heat_map=True,
):
    """
    Visualize and add tensors of n-dimentionals to a Tensorboard SummaryWriter. Tensors
    will be visualized as a 2D grid image.
    Args:
        writer (SummaryWriter): Tensorboard SummaryWriter.
        array (tensor): tensor to visualize.
        name (str): name of the tensor.
        nrow (Optional[int]): number of 2D filters in each row in the grid image.
        normalize (bool): whether to normalize when we have multiple 2D filters.
            Default to False.
        global_step (Optional[int]): current step.
        heat_map (bool): whether to add heat map to 2D each 2D filters in array.
    """
    if array is not None and array.ndim != 0:
        if array.ndim == 1:
            reshaped_array = array.unsqueeze(0)
            if nrow is None:
                nrow = int(math.sqrt(reshaped_array.size()[1]))
            reshaped_array = reshaped_array.view(-1, nrow)
            if heat_map:
                reshaped_array = add_heatmap(reshaped_array)
                writer.add_image(
                    name,
                    reshaped_array,
                    global_step=global_step,
                    dataformats="CHW",
                )
            else:
                writer.add_image(
                    name,
                    reshaped_array,
                    global_step=global_step,
                    dataformats="HW",
                )
        elif array.ndim == 2:
            reshaped_array = array
            if heat_map:
                heatmap = add_heatmap(reshaped_array)
                writer.add_image(
                    name, heatmap, global_step=global_step, dataformats="CHW"
                )
            else:
                writer.add_image(
                    name,
                    reshaped_array,
                    global_step=global_step,
                    dataformats="HW",
                )
        else: