Python Package

pytorch-graph

A lightweight PyTorch utility for visualising computational graphs and architecture diagrams. Use this page as an exploration surface for quick starts, common examples, and project links.

PyPI Documentation GitHub Issues

Feature highlights

Architecture visualisation

Generate vertical flowchart and publication-ready research-paper styles.
Inspect full layer names, tensor dimensions, and parameter counts.
Export high-DPI PNG outputs suitable for notes, reports, and papers.

Computational graph analysis

Capture complete autograd graphs with optional memory and timing tracking.
Inspect execution flow, tensor operations, and graph-level metrics.
Keep full operation names visible to avoid truncated debugging output.

Model complexity insights

Get rapid summaries of parameter counts and total size.
Track memory estimates per layer and across the full model.
Switch between compact and detail-rich outputs for different use cases.

Quick install

Base install

pip install pytorch-graph

Extended features

pip install pytorch-graph[full]

Development extra

pip install pytorch-graph[dev]

Exploration examples

Architecture diagramimport torch
import torch.nn as nn
from pytorch_graph import generate_architecture_diagram

model = nn.Sequential(nn.Linear(784, 128), nn.ReLU(), nn.Linear(128, 10))
generate_architecture_diagram(
  model=model,
  input_shape=(1, 784),
  output_path="model_architecture.png",
  title="MNIST MLP",
  style="flowchart",
)

Full graph trackingfrom pytorch_graph import ComputationalGraphTracker
import torch

tracker = ComputationalGraphTracker(model=model, track_memory=True, track_timing=True)
tracker.start_tracking()
output = model(torch.randn(1, 784))
tracker.stop_tracking()
tracker.save_graph_png("complete_graph.png", width=1800, height=1200, dpi=300)

Model analysisfrom pytorch_graph import analyze_model

analysis = analyze_model(model=model, input_shape=(1, 784), detailed=True)
print(analysis["summary"]["total_params"])
print(analysis["summary"]["trainable_params"])