LeNet Style Without Funnel Connectors

Compare the default tapered connectors with plain box-to-box connectors by setting draw_funnel=False.

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

# Example of a simple CNN model using nn.Sequential
model = nn.Sequential(
    nn.Conv2d(3, 8, kernel_size=3, padding=1),
    nn.MaxPool2d(2, 2),
    nn.Conv2d(8, 16, kernel_size=3, padding=1),
    nn.MaxPool2d(2, 2),
)

input_shape = (1, 3, 128, 128)

img_default = visualtorch.render(model, input_shape=input_shape, style="lenet")
img_no_funnel = visualtorch.render(model, input_shape=input_shape, style="lenet", draw_funnel=False)

dpi = 150  # rendered at 2x this in the final doc build (savefig.dpi=300 in conf.py)
_, axes = plt.subplots(
    1,
    2,
    figsize=((img_default.width + img_no_funnel.width) / dpi, max(img_default.height, img_no_funnel.height) / dpi),
    dpi=dpi,
)

for ax, img, title in zip(axes, [img_default, img_no_funnel], ["Default", "draw_funnel=False"], strict=True):
    ax.imshow(img)
    ax.set_title(title)
    ax.axis("off")

plt.tight_layout()
plt.show()