Legend Position

Legend Position#

Place the graph legend above or below the diagram, aligned left, center, or right.

import matplotlib.pyplot as plt
import visualtorch
from torch import nn

model = nn.Sequential(
    nn.Conv2d(3, 8, kernel_size=3, padding=1),
    nn.ReLU(),
    nn.Conv2d(8, 16, kernel_size=3, padding=1),
    nn.ReLU(),
)

input_shape = (1, 3, 32, 32)
dpi = 150  # rendered at 2x this in the final doc build (savefig.dpi=300 in conf.py)


def _show(position: str) -> None:
    img = visualtorch.render(
        model,
        input_shape=input_shape,
        style="graph",
        show_neurons=False,
        legend=True,
        legend_position=position,
        layer_spacing=400,
    )
    plt.figure(figsize=(img.width / dpi, img.height / dpi), dpi=dpi)
    plt.imshow(img)
    plt.axis("off")
    plt.tight_layout()
    plt.show()