u/LostDistance9365

I built a small tool to auto-generate architecture diagrams directly from PyTorch models, which I originally built for my own research paper.

26k+ PyPI downloads, already used in publications (Nature, IEEE, MDPI), check out some use cases here: https://visualtorch.readthedocs.io/en/latest/markdown/showcase/index.html

It traces an actual forward pass, so it correctly captures branching, skip connections, and multi-input models, not just flat sequential stacks.

import visualtorch
import torchvision.models as models

model = models.resnet18()
img = visualtorch.render(model, input_shape=(1, 3, 224, 224), style="graph", show_neurons=False, layer_spacing=60)
img.save("resnet18.png")

Three rendering styles depending on what you want to show:

• graph: node/edge diagram, good for showing branching/skip connections clearly
• flow: stacked volumetric boxes, closer to the classic CNN-paper look
• lenet: the classic LeNet stacked-plane style

GitHub: https://github.com/willyfh/visualtorch | Docs: https://visualtorch.readthedocs.io/en/latest/

Open to feedback, especially if you hit a model it renders weirdly :)