Function plot_conv2d_weights

bindsnet/analysis/plotting.py:264–319

Plot a connection weight matrix of a Conv2dConnection. :param weights: Weight matrix of Conv2dConnection object. :param wmin: Minimum allowed weight value. :param wmax: Maximum allowed weight value. :param im: Used for re-drawing the weights plot. :param figsize: Horizontal

(
    weights: torch.Tensor,
    wmin: float = 0.0,
    wmax: float = 1.0,
    im: Optional[AxesImage] = None,
    figsize: Tuple[int, int] = (5, 5),
    cmap: str = "hot_r",
    title: Optional[str] = None,
)

def plot_conv2d_weights(
    weights: torch.Tensor,
    wmin: float = 0.0,
    wmax: float = 1.0,
    im: Optional[AxesImage] = None,
    figsize: Tuple[int, int] = (5, 5),
    cmap: str = "hot_r",
    title: Optional[str] = None,
) -> AxesImage:
    # language=rst
    """
    Plot a connection weight matrix of a Conv2dConnection.

    :param weights: Weight matrix of Conv2dConnection object.
    :param wmin: Minimum allowed weight value.
    :param wmax: Maximum allowed weight value.
    :param im: Used for re-drawing the weights plot.
    :param figsize: Horizontal, vertical figure size in inches.
    :param cmap: Matplotlib colormap.
    :param title: Title of the plot.
    :return: Used for re-drawing the weights plot.
    """

    sqrt1 = int(np.ceil(np.sqrt(weights.size(0))))
    sqrt2 = int(np.ceil(np.sqrt(weights.size(1))))
    height, width = weights.size(2), weights.size(3)
    reshaped = reshape_conv2d_weights(weights)

    if not im:
        fig, ax = plt.subplots(figsize=figsize)
        im = ax.imshow(reshaped, cmap=cmap, vmin=wmin, vmax=wmax, aspect="auto")
        div = make_axes_locatable(ax)
        cax = div.append_axes("right", size="5%", pad=0.05)

        for i in range(height, sqrt1 * sqrt2 * height, height):
            ax.axhline(i - 0.5, color="g", linestyle="--")
            if i % sqrt1 == 0:
                ax.axhline(i - 0.5, color="g", linestyle="-")

        for i in range(width, sqrt1 * sqrt2 * width, width):
            ax.axvline(i - 0.5, color="g", linestyle="--")
            if i % sqrt1 == 0:
                ax.axvline(i - 0.5, color="g", linestyle="-")

        ax.set_xticks(())
        ax.set_yticks(())
        ax.set_aspect("auto")
        if title != None:
            ax.set_title(title + " Weights")

        plt.colorbar(im, cax=cax)
        fig.tight_layout()
    else:
        im.set_data(reshaped)

    return im