Neural Network Architecture Diagram

A neural network architecture diagram shows the layers of a model in order and the shape of the tensor flowing between them. This sample is a small convolutional network that classifies 28x28 handwritten digits (MNIST): two convolution and pooling stages extract features, then a dense head picks one of ten classes. These are deep-learning layer diagrams, not UML; UML covers software structure and behaviour, and layer diagrams are the standard notation used alongside it for model internals.

1. The Diagram

Small CNN for MNIST Digits: Layers and Tensor Shapes Input 1x28x28 grayscale 0 Conv1 32x28x28 3x3, pad 1, ReLU 320 Pool1 32x14x14 max 2x2 0 Conv2 64x14x14 3x3, pad 1, ReLU 18,496 Pool2 64x7x7 max 2x2 0 Flatten 3136 64*7*7 0 Dense 128 ReLU, dropout 401,536 Output 10 softmax 1,290 shape (C x H x W) parameters Total 421,642 trainable parameters; 95% of them sit in the first Dense layer.

Box height is drawn roughly in proportion to the size of each layer's output, so you can see the spatial maps shrink while the channel count grows.

  1. Input: one grayscale channel, 28x28 pixels, so the tensor is 1x28x28 (PyTorch's channels-first order).
  2. Conv1: 32 filters of 3x3 with padding 1 keep the 28x28 size and produce 32 feature maps. Parameters: 3*3*1*32 weights + 32 biases = 320.
  3. Pool1: 2x2 max pooling halves height and width to 14x14 and has no parameters.
  4. Conv2: 64 filters of 3x3 over 32 input channels: 3*3*32*64 + 64 = 18,496 parameters.
  5. Pool2 shrinks to 64x7x7, and Flatten turns it into a vector of 3,136 numbers.
  6. Dense: 3,136*128 + 128 = 401,536 parameters, by far the largest layer, followed by dropout to fight overfitting.
  7. Output: 10 units, one per digit; softmax turns the scores into probabilities.

2. How to Read It

Notation Used Layer box height ~ size of the activation it outputs Input Output 1x28x28 output shape, channels first 18,496 learnable weights + biases Arrows are tensors flowing forward; the batch dimension N is omitted.

Read left to right as the forward pass. Shapes under each box are the layer's output; checking that each shape follows from the previous one is the fastest way to catch a wrong padding or stride before running any code.

3. Variant: Neuron Graph

Variant: The Dense Head as a Neuron Graph (MLP 784-128-10) ... Input 784 flattened pixels ... Hidden 128 ReLU ... Output 10 softmax W1: 784x128 W2: 128x10

The classic neuron-and-edge drawing shows a plain multilayer perceptron, where every unit connects to every unit in the next layer. It explains what "fully connected" means, but it cannot show 784 inputs, so ellipses stand in for the missing nodes. For anything bigger than a toy network, prefer the layer-box style above.

4. Pros, Cons, and When to Use

Pros

Cons

When to use

When not to use

5. As Diagrams-as-Code

The layer chain in Graphviz DOT (record nodes carry the shapes):

digraph cnn {
  rankdir=LR; node [shape=record, style=rounded];
  in    [label="Input|1x28x28"];
  c1    [label="Conv1 3x3 32, ReLU|32x28x28|320"];
  p1    [label="MaxPool 2x2|32x14x14"];
  c2    [label="Conv2 3x3 64, ReLU|64x14x14|18,496"];
  p2    [label="MaxPool 2x2|64x7x7"];
  fl    [label="Flatten|3136"];
  d1    [label="Dense 128, ReLU|128|401,536"];
  out   [label="Dense 10, softmax|10|1,290"];
  in -> c1 -> p1 -> c2 -> p2 -> fl -> d1 -> out;
}