Jesse Roberts

Jesse Roberts

Native Tennesseean Researching AI, Language, and Robotics

Torch Canvas

Diagram ⇄ PyTorch

A visual editor for neural networks that writes PyTorch as you draw, and turns PyTorch modules back into diagrams. It runs in the browser with nothing to install.

Open Torch Canvas Source on GitHub

Start from a real model

Each one opens in the editor, ready to explore, change, and export. 22 networks and 26 building blocks, straight from the project's example library. A few favorites first, then everything by category.

Diagram of GPT-2 (124M architecture) GPT-2 (124M architecture) 124.44M GPT-2 small: token and position embeddings, 12 × GPT2Block (pre-norm causal self-attention and a tanh-GELU MLP, with dropout), a final LayerNorm and the LM head. GPT-2... Uses GPT2Block Open in Torch Canvas → Diagram of Llama 2 7B Llama 2 7B 6.74B Llama 2 7B as LlamaForCausalLM: token embedding, 32 × LlamaDecoderLayer (pre-norm with RMSNorm; causal self-attention with rotary position embeddings on q and k; a Swi... Uses LlamaDecoderLayer Open in Torch Canvas → Diagram of BERT-base (masked LM) BERT-base (masked LM) 109.51M BERT-base as BertForMaskedLM: word, position and token-type embeddings with a LayerNorm, 12 × BertLayer (post-norm: bidirectional self-attention, then a GELU feed-forw... Uses BertLayer Open in Torch Canvas → Diagram of Qwen2.5-VL 3B Qwen2.5-VL 3B 3.75B Qwen2.5-VL-3B-Instruct as Qwen2_5_VLForConditionalGeneration, giving the same output as the reference implementation in Hugging Face transformers. Inputs as Qwen's pro... Uses Qwen2_5_VLVisionBlock, Qwen2_5_VLVisionStage, Qwen2_5_VLPatchMerger, Qwen2_5_VLDecoderLayer Open in Torch Canvas →

Browse all 48 examples

MLPs and residual basics

Small fully connected networks for seeing how depth, activations, normalization and residuals behave.

Diagram of Residual MLP Residual MLP 1.52M The reference diagram: a down projection, three gated residual branches and a gated head, every layer drawn out. Open in Torch Canvas →
Diagram of Residual MLP with a block Residual MLP with a block 1.52M The same network with the gated branch defined once as GatedBlock and repeated 3 times. Uses GatedBlock Open in Torch Canvas →
Diagram of MLP, 10 layers, ReLU MLP, 10 layers, ReLU 729.9K A plain deep MLP: 10 Linear layers with ReLU between them. Good for seeing how depth and init affect the signal. Open in Torch Canvas →
Diagram of MLP, 10 layers, Tanh MLP, 10 layers, Tanh 729.9K The same 10-layer MLP with Tanh between the layers. Open in Torch Canvas →
Diagram of Pre-norm residual MLP Pre-norm residual MLP 2.31M A transformer-style stack: four pre-LayerNorm GELU MLP blocks with residual connections, then a classifier. Open in Torch Canvas →

Building blocks

Diagram of GatedBlock GatedBlock 393.5K Open in Torch Canvas →
Diagram of LinearReLU LinearReLU 65.8K Open in Torch Canvas →
Diagram of PreNormMLP PreNormMLP 526.1K Open in Torch Canvas →
Diagram of SwiGLU SwiGLU 393.2K Open in Torch Canvas →

Convolutional networks

Image classification, detection and segmentation built from convolutions.

Diagram of Basic CNN (MNIST) Basic CNN (MNIST) 421.6K A plain convolutional classifier for 28 × 28 digits, no blocks: two Conv 3×3 → ReLU → max pool stages (32 and 64 channels), then flatten, Linear 3136 → 128 → ReLU → dr... Open in Torch Canvas →
Diagram of ResNet-50 (ImageNet) ResNet-50 (ImageNet) 25.56M ResNet-50 as in torchvision (v1.5, stride on the 3×3): a 7×7 stem and max pool, then four stages of bottleneck blocks (3, 4, 6, 3) at 256 / 512 / 1024 / 2048 channels,... Uses Bottleneck, BottleneckDown Open in Torch Canvas →
Diagram of YOLOv3 (416, COCO) YOLOv3 (416, COCO) 61.95M YOLOv3 as in Darknet's yolov3.cfg: the Darknet-53 backbone (Conv-BN-LeakyReLU layers and 1 + 2 + 8 + 8 + 4 residual blocks), then detection at three scales: 13 × 13 fo... Uses DarknetConv, DarknetResidual, YoloConvSet, YoloDetect Open in Torch Canvas →
Diagram of U-Net (segmentation) U-Net (segmentation) 31.04M U-Net as in milesial/Pytorch-UNet (bilinear=False): an encoder of DoubleConv and four Down steps (64 → 1024 channels, 256 → 16 pixels), a decoder of four Up steps that... Uses DoubleConv, Down, Up Open in Torch Canvas →

Building blocks

Diagram of DarknetConv DarknetConv 18.6K Open in Torch Canvas →
Diagram of DarknetResidual DarknetResidual 20.7K Open in Torch Canvas →
Diagram of Bottleneck Bottleneck 70.4K Open in Torch Canvas →
Diagram of BottleneckDown BottleneckDown 379.4K Open in Torch Canvas →
Diagram of DoubleConv DoubleConv 221.7K Open in Torch Canvas →
Diagram of Down Down 221.7K Open in Torch Canvas →
Diagram of Up Up 143.7K Open in Torch Canvas →
Diagram of YoloConvSet YoloConvSet 11.02M Open in Torch Canvas →
Diagram of YoloDetect YoloDetect 4.98M Open in Torch Canvas →

Recurrent and state-space models

Networks that carry a state through a sequence, from the Elman RNN to Mamba.

Diagram of RNN (Elman) RNN (Elman) 5.8K A basic recurrent network drawn flat: h = tanh(W_ih·x + W_hh·h_prev), where h_prev comes from a Delay (the previous time step, zeros at the start). The code runs it as... Open in Torch Canvas →
Diagram of LSTM LSTM 21.5K A long short-term memory network drawn flat: four gates (input i, forget f, cell g, output o), each from x and from h_prev; then c = f ⊙ c_prev + i ⊙ g and h = o ⊙ tan... Open in Torch Canvas →
Diagram of GRU GRU 16.3K A gated recurrent unit drawn flat: reset r and update z gates from x and h_prev, the candidate n = tanh(W_in·x + r ⊙ (W_hn·h_prev)), then h = (1 − z) ⊙ n + z ⊙ h_prev,... Open in Torch Canvas →
Diagram of Mamba-130M Mamba-130M 129.14M Mamba (state-spaces/mamba-130m, as transformers' MambaForCausalLM): token embeddings, 24 MambaBlocks (RMSNorm, then the selective state space mixer: input projections,... Uses MambaBlock Open in Torch Canvas →

Building blocks

Diagram of MambaBlock MambaBlock 3.77M Open in Torch Canvas →

Transformers and language models

The original encoder-decoder Transformer and the language models built on it.

Diagram of Decoder-only transformer, fully expanded Decoder-only transformer, fully expanded 92.16M A 2-layer GPT-style decoder with every layer drawn out: token and position embeddings, causal multi-head self-attention, a GELU MLP, residuals, LayerNorms and the LM h... Open in Torch Canvas →
Diagram of GPT-2 (124M architecture) GPT-2 (124M architecture) 124.44M GPT-2 small: token and position embeddings, 12 × GPT2Block (pre-norm causal self-attention and a tanh-GELU MLP, with dropout), a final LayerNorm and the LM head. GPT-2... Uses GPT2Block Open in Torch Canvas →
Diagram of BERT-base (masked LM) BERT-base (masked LM) 109.51M BERT-base as BertForMaskedLM: word, position and token-type embeddings with a LayerNorm, 12 × BertLayer (post-norm: bidirectional self-attention, then a GELU feed-forw... Uses BertLayer Open in Torch Canvas →
Diagram of Llama 2 7B Llama 2 7B 6.74B Llama 2 7B as LlamaForCausalLM: token embedding, 32 × LlamaDecoderLayer (pre-norm with RMSNorm; causal self-attention with rotary position embeddings on q and k; a Swi... Uses LlamaDecoderLayer Open in Torch Canvas →
Diagram of Transformer (encoder-decoder) Transformer (encoder-decoder) 63.08M The original Transformer, "Attention Is All You Need" (base: d = 512, 8 heads, feed-forward 2048, 6 + 6 layers, dropout 0.1), computed as nn.Transformer does (post-nor... Uses EncoderLayer, DecoderLayer Open in Torch Canvas →

Building blocks

Diagram of BertLayer BertLayer 7.09M Open in Torch Canvas →
Diagram of GPT2Block GPT2Block 7.09M Open in Torch Canvas →
Diagram of LlamaDecoderLayer LlamaDecoderLayer 202.38M Open in Torch Canvas →
Diagram of DecoderLayer DecoderLayer 4.20M Open in Torch Canvas →
Diagram of EncoderLayer EncoderLayer 3.15M Open in Torch Canvas →

Vision transformers and multimodal

Transformers that read images, alone or together with text.

Diagram of Qwen2.5-VL 3B Qwen2.5-VL 3B 3.75B Qwen2.5-VL-3B-Instruct as Qwen2_5_VLForConditionalGeneration, giving the same output as the reference implementation in Hugging Face transformers. Inputs as Qwen's pro... Uses Qwen2_5_VLVisionBlock, Qwen2_5_VLVisionStage, Qwen2_5_VLPatchMerger, Qwen2_5_VLDecoderLayer Open in Torch Canvas →
Diagram of ViT-B/16 (ImageNet) ViT-B/16 (ImageNet) 86.57M The Vision Transformer, as torchvision's vit_b_16: the 224 × 224 image is cut into 16 × 16 patches by a strided convolution (196 patches of 768 features), a learned cl... Uses ViTEncoderBlock Open in Torch Canvas →

Building blocks

Diagram of Qwen2_5_VLDecoderLayer Qwen2_5_VLDecoderLayer 77.08M Open in Torch Canvas →
Diagram of Qwen2_5_VLPatchMerger Qwen2_5_VLPatchMerger 36.71M Open in Torch Canvas →
Diagram of Qwen2_5_VLVisionBlock Qwen2_5_VLVisionBlock 19.70M Open in Torch Canvas →
Diagram of Qwen2_5_VLVisionStage Qwen2_5_VLVisionStage 157.62M Open in Torch Canvas →
Diagram of ViTEncoderBlock ViTEncoderBlock 7.09M Open in Torch Canvas →

Generative and self-supervised

Models that learn without labels, by reconstructing their input or predicting hidden parts of it.

Diagram of VAE (MNIST) VAE (MNIST) 652.8K A variational autoencoder, as in PyTorch's examples/vae: the encoder (784 → 400, ReLU) gives the mean μ and log-variance of a 20-dimensional Gaussian; z = μ + ε·exp(½·... Open in Torch Canvas →
Diagram of I-JEPA (small) I-JEPA (small) 5.75M params · ❄ 2.70M A small I-JEPA (facebookresearch/ijepa: ViT encoder and predictor; 32 × 32 images, 4 × 4 patches, width 192, 6 + 3 layers). The context encoder sees only the context p... Uses JEPABlock, Encoder Open in Torch Canvas →

Building blocks

Diagram of JEPABlock JEPABlock 444.9K Open in Torch Canvas →
Diagram of Encoder Encoder 2.67M Open in Torch Canvas →

What it does

Diagram ⇄ code, both ways

Draw a network and the PyTorch code updates as you go. Paste an nn.Module (or a torch.fx export) and it becomes a diagram. Every example round-trips: diagram → code → the same diagram.

Shapes and parameters, live

Tensor sizes are worked out symbolically and shown on every connection, with a running parameter count. A size that doesn't fit, like a head count that doesn't divide d, turns red.

See variance before you train

Plot how activation and gradient variance change through the network right after initialization, worked out analytically. Change an init and replot to compare.

Reusable blocks

Define a block once (a transformer layer, a SwiGLU MLP) and repeat it. Edit it anywhere and every network that uses it updates.

Ready-to-run training scripts

Add a loss (cross-entropy with next-token or masked targets, BCE, MSE, L1, Huber), then export model.py and a complete train.py with your optimizer, schedule, precision, and checkpointing.

Real architectures, checked

GPT-2, BERT, Llama 2 and Qwen2.5-VL match their reference implementations: the same parameter counts at full size, and the same outputs with the weights copied across (GPT-2 against nanoGPT; the others against Hugging Face transformers).

For teaching

Every example has its own link, so a slide or course page can open exactly the model being discussed: https://www.jessetnroberts.com/torch-canvas/?open=gpt2 opens GPT-2. It's used in AI 4200 / 5200 Deep Learning.

Example list synced from JesseTNRoberts/torch-canvas (b95b148) on October 23, 2026.