Networks¶
Complete networks written in HNDL, from examples/networks/. Each entry is resolved from its
declared contract; the shape tables and parameter counts below are generated by python -m hndl.docs.
Fully connected autoencoder¶
Compress a 28×28 image through a 32-unit bottleneck and reconstruct it. The decoder's final width is inferred from the reshape target and the output contract.
examples/networks/autoencoder.hndl
# A fully connected autoencoder for 28×28 grayscale images (MNIST-sized).
# The image is flattened to 784 values, squeezed through a 32-unit bottleneck,
# and reconstructed back into the same [B, 1, 28, 28] shape.
# Encoder: 784 -> 128 -> 32.
flatten()
linear(128)
relu()
linear(32, name="code") # The bottleneck; model["code"] is the 32-unit layer.
relu()
# Decoder: 32 -> 128 -> 784. The final width is inferred from the reshape
# target, which the output contract fixes at 1·28·28 = 784.
linear(128)
relu()
linear(name="reconstruction")
sigmoid() # Pixels in (0, 1), matching normalized inputs.
reshape(1, 28, 28)
Input ['B', 1, 28, 28] → output ['B', 1, 28, 28].
Network: [B, 1, 28, 28] -> [B, 1, 28, 28] dtype=float32
index name operation input shapes output shapes
0 n0 flatten x=[B, 1, 28, 28] out=[B, 784]
1 n1 linear x=[B, 784] out=[B, 128]
2 n2 relu x=[B, 128] out=[B, 128]
3 code linear x=[B, 128] out=[B, 32]
4 n4 relu x=[B, 32] out=[B, 32]
5 n5 linear x=[B, 32] out=[B, 128]
6 n6 relu x=[B, 128] out=[B, 128]
7 reconstruction linear x=[B, 128] out=[B, 784]
8 n8 sigmoid x=[B, 784] out=[B, 784]
9 n9 reshape x=[B, 784] out=[B, 1, 28, 28]
Parameters: 209,968. Classic MNIST autoencoder with a 784-128-32-128-784 stack; 784·128+128 + 128·32+32 + 32·128+128 + 128·784+784 = 209,968 parameters, cross-checked against the equivalent torch.nn.Sequential.
Conditional GAN discriminator¶
A conditional discriminator with two named inputs and two named outputs. The one-hot label y is projected onto a 28×28 plane and stacked on the image x as a second channel; two "down2" convolutions reduce the result to a 128×7×7 map. The flattened map is published as features for a feature-matching loss, and the head produces raw logits — apply a sigmoid, or use BCEWithLogitsLoss, in the loss rather than in the network.
examples/networks/conditional_discriminator.hndl
# A conditional GAN discriminator for 28x28 grayscale digits. It takes two
# external inputs: the image "x" and a one-hot class label "y".
# Project the label onto its own 28x28 plane and stack it on the image as an
# extra channel, the usual conditional-GAN trick.
plane = linear(y, 784, name="label_projection")
label = reshape(plane, 1, 28, 28, name="label_plane")
concat(x, label, name="conditioned") # [B, 2, 28, 28]
# Two strided convolutions halve the image: 28 -> 14 -> 7.
conv(64, policy="down2", name="stage1")
leaky_relu(0.2)
conv(128, policy="down2", bias=False, name="stage2")
batch_norm()
leaky_relu(0.2)
# Both public outputs are named: the 128x7x7 map flattened into "features",
# which a feature-matching loss can read, and the real/fake "logits".
features = flatten(name="features")
logits = linear(1, name="logits")
Input x=['B', 1, 28, 28], y=['B', 10] → output logits=['B', 1], features=['B', 6272].
Network: x=[B, 1, 28, 28], y=[B, 10] -> logits=[B, 1], features=[B, 6272] dtype=float32
index name operation input shapes output shapes
0 label_projection linear x=[B, 10] out=[B, 784]
1 label_plane reshape x=[B, 784] out=[B, 1, 28, 28]
2 conditioned concat x0=[B, 1, 28, 28], x1=[B, 1, 28, 28] out=[B, 2, 28, 28]
3 stage1 conv x=[B, 2, 28, 28] out=[B, 64, 14, 14]
4 n4 leaky_relu x=[B, 64, 14, 14] out=[B, 64, 14, 14]
5 stage2 conv x=[B, 64, 14, 14] out=[B, 128, 7, 7]
6 n6 batch_norm x=[B, 128, 7, 7] out=[B, 128, 7, 7]
7 n7 leaky_relu x=[B, 128, 7, 7] out=[B, 128, 7, 7]
8 features flatten x=[B, 128, 7, 7] out=[B, 6272]
9 logits linear x=[B, 6272] out=[B, 1]
Parameters: 148,337. Mirza & Osindero, "Conditional Generative Adversarial Nets" (2014), in the usual DCGAN form; 10·784+784 + 2·64·4·4+64 + 64·128·4·4 + 2·128 + 6272+1 = 148,337 parameters.
DCGAN discriminator¶
The DCGAN discriminator for 64×64 RGB images: four "down2" convolutions with leaky ReLU, batch-normalized after the first stage, reduce the image to a 512×4×4 map that is flattened and projected to a single logit. The output is a raw score — apply a sigmoid, or use BCEWithLogitsLoss, in the loss rather than in the network.
examples/networks/dcgan_discriminator.hndl
# DCGAN discriminator (Radford et al., 2015) for 64x64 RGB images.
# Four strided convolutions halve the image each time: 64 -> 32 -> 16 -> 8 -> 4.
# The "down2" policy supplies the canonical kernel 4, stride 2, padding 1.
# The first stage has no normalization, so it keeps its bias.
conv(64, policy="down2", name="stage1")
leaky_relu(0.2)
# Every later stage is conv -> batch_norm -> leaky_relu. batch_norm's own bias
# makes the convolution bias redundant, so bias=False there.
conv(128, policy="down2", bias=False, name="stage2")
batch_norm()
leaky_relu(0.2)
conv(256, policy="down2", bias=False, name="stage3")
batch_norm()
leaky_relu(0.2)
conv(512, policy="down2", bias=False, name="stage4")
batch_norm()
leaky_relu(0.2)
# The 512x4x4 map flattens to 8192; the head's width follows the output contract.
flatten()
linear(name="logit")
Input ['B', 3, 64, 64] → output ['B', 1].
Network: [B, 3, 64, 64] -> [B, 1] dtype=float32
index name operation input shapes output shapes
0 stage1 conv x=[B, 3, 64, 64] out=[B, 64, 32, 32]
1 n1 leaky_relu x=[B, 64, 32, 32] out=[B, 64, 32, 32]
2 stage2 conv x=[B, 64, 32, 32] out=[B, 128, 16, 16]
3 n3 batch_norm x=[B, 128, 16, 16] out=[B, 128, 16, 16]
4 n4 leaky_relu x=[B, 128, 16, 16] out=[B, 128, 16, 16]
5 stage3 conv x=[B, 128, 16, 16] out=[B, 256, 8, 8]
6 n6 batch_norm x=[B, 256, 8, 8] out=[B, 256, 8, 8]
7 n7 leaky_relu x=[B, 256, 8, 8] out=[B, 256, 8, 8]
8 stage4 conv x=[B, 256, 8, 8] out=[B, 512, 4, 4]
9 n9 batch_norm x=[B, 512, 4, 4] out=[B, 512, 4, 4]
10 n10 leaky_relu x=[B, 512, 4, 4] out=[B, 512, 4, 4]
11 n11 flatten x=[B, 512, 4, 4] out=[B, 8192]
12 logit linear x=[B, 8192] out=[B, 1]
Parameters: 2,765,633. Radford, Metz & Chintala, "Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks" (2015), in the form of the PyTorch DCGAN example; the count was checked against the equivalent torch.nn stack — 3·64·4·4+64 + 64·128·4·4 + 2·128 + 128·256·4·4 + 2·256 + 256·512·4·4 + 2·512 + 8192+1 = 2,765,633.
DCGAN generator¶
Project a 100-dimensional latent vector to a 4×4×1024 seed, then upsample with four doubling transposed convolutions to a 3×64×64 image. The projection width, the seed size and the three output channels are all inferred from the contract.
examples/networks/dcgan_generator.hndl
# DCGAN generator (Radford et al. 2015, figure 1): a 100-dimensional latent
# vector becomes a 64×64 RGB image. Four doubling stages divide 64 by 16, so
# the contract fixes the seed at 4×4 and the projection at 1024·4·4 = 16384.
# Every convolution that feeds a batch_norm drops its bias, as in the reference.
linear(bias=False, name="project") # [B, 100] -> [B, 16384]
reshape(1024, name="seed") # -> [B, 1024, 4, 4]
batch_norm(name="seed_norm")
relu()
deconv(512, policy="up2", bias=False, name="up8") # 4×4 -> 8×8
batch_norm()
relu()
deconv(256, policy="up2", bias=False, name="up16") # 8×8 -> 16×16
batch_norm()
relu()
deconv(128, policy="up2", bias=False, name="up32") # 16×16 -> 32×32
batch_norm()
relu()
# 32×32 -> 64×64. The 3 output channels come from the contract, and this last
# stage keeps its bias because no normalization follows it.
deconv(policy="up2", name="to_rgb")
tanh() # images in (-1, 1)
Input ['B', 100] → output ['B', 3, 64, 64].
Network: [B, 100] -> [B, 3, 64, 64] dtype=float32
index name operation input shapes output shapes
0 project linear x=[B, 100] out=[B, 16384]
1 seed reshape x=[B, 16384] out=[B, 1024, 4, 4]
2 seed_norm batch_norm x=[B, 1024, 4, 4] out=[B, 1024, 4, 4]
3 n3 relu x=[B, 1024, 4, 4] out=[B, 1024, 4, 4]
4 up8 deconv x=[B, 1024, 4, 4] out=[B, 512, 8, 8]
5 n5 batch_norm x=[B, 512, 8, 8] out=[B, 512, 8, 8]
6 n6 relu x=[B, 512, 8, 8] out=[B, 512, 8, 8]
7 up16 deconv x=[B, 512, 8, 8] out=[B, 256, 16, 16]
8 n8 batch_norm x=[B, 256, 16, 16] out=[B, 256, 16, 16]
9 n9 relu x=[B, 256, 16, 16] out=[B, 256, 16, 16]
10 up32 deconv x=[B, 256, 16, 16] out=[B, 128, 32, 32]
11 n11 batch_norm x=[B, 128, 32, 32] out=[B, 128, 32, 32]
12 n12 relu x=[B, 128, 32, 32] out=[B, 128, 32, 32]
13 to_rgb deconv x=[B, 128, 32, 32] out=[B, 3, 64, 64]
14 n14 tanh x=[B, 3, 64, 64] out=[B, 3, 64, 64]
Parameters: 12,658,435. Radford, Metz & Chintala 2015 (DCGAN), figure 1: project and reshape to 4×4×1024, then 8×8×512, 16×16×256, 32×32×128, 64×64×3. The count was checked by building the same stack in plain torch — Linear(100, 16384, bias=False), four ConvTranspose2d(kernel 4, stride 2, padding 1) with bias=False under each of the four BatchNorm2d layers and bias=True on the final RGB stage — which also gives 12,658,435. The PyTorch DCGAN tutorial differs: it starts from a 100×1×1 ConvTranspose2d rather than a linear projection, and its ngf=64 feature maps make the stack 512/256/128/64/3 instead of the paper's 1024/512/256/128/3.
Tiny GPT¶
A 64-token GPT-2-style decoder: learned token and position embeddings, four causal pre-norm transformer blocks with a tanh-approximated GELU feed-forward, a final layer norm, and an untied language-model head whose width is inferred from the output contract.
examples/networks/gpt_tiny.hndl
# A tiny GPT-2-style decoder: 64 token ids in, one logit per vocabulary entry out.
# Vocabulary 256, model width 128, four causal blocks.
# Stem: a vector per token id, plus a learned absolute position for each of the 64 slots.
embedding(256, 128, name="wte")
pos_embed(64, name="wpe")
# Four pre-norm causal blocks: 4 heads of width 32, 4x feed-forward, tanh-approximated GELU.
transformer_block(4, activation="gelu_tanh", causal=True, name="block0")
transformer_block(4, activation="gelu_tanh", causal=True, name="block1")
transformer_block(4, activation="gelu_tanh", causal=True, name="block2")
transformer_block(4, activation="gelu_tanh", causal=True, name="block3")
# Final normalization, then the language-model head. Its width is the vocabulary
# size, so the output contract determines it; GPT-2 ties it to wte, HNDL does not.
layer_norm(name="ln_f")
linear(bias=False, name="lm_head")
Input ['B', 64] (input_dtype="int64") → output ['B', 64, 256].
Network: [B, 64] -> [B, 64, 256] dtype=float32 input_dtype=int64
index name operation input shapes output shapes
0 wte embedding ids=[B, 64] out=[B, 64, 128]
1 wpe pos_embed x=[B, 64, 128] out=[B, 64, 128]
2 block0 transformer_block x=[B, 64, 128] out=[B, 64, 128]
3 block1 transformer_block x=[B, 64, 128] out=[B, 64, 128]
4 block2 transformer_block x=[B, 64, 128] out=[B, 64, 128]
5 block3 transformer_block x=[B, 64, 128] out=[B, 64, 128]
6 ln_f layer_norm x=[B, 64, 128] out=[B, 64, 128]
7 lm_head linear x=[B, 64, 128] out=[B, 64, 256]
Parameters: 867,072. Block structure follows Radford et al. 2019, "Language Models are Unsupervised Multitask Learners" (GPT-2); the count was checked against the same architecture built in plain torch (256·128 + 64·128 + 4·(4·128 + 4·(128·128+128) + (512·128+512) + (128·512+128)) + 2·128 + 128·256 = 867,072), with the head untied rather than shared with the token embedding.
Hopfield classifier¶
Classify a 28×28 image by projecting it to a 128-wide query, retrieving from a 64-pattern modern Hopfield memory, and reading out the classes. The final width is inferred from the output contract.
examples/networks/hopfield_classifier.hndl
# A modern Hopfield layer as the hidden layer of a 28×28 grayscale classifier.
flatten()
# Project the image to the query width; the memory acts on this last axis.
linear(128, name="query")
# 64 stored patterns of width 128, the only parameters of the layer.
# beta=0.25 keeps the softmax soft, so a query retrieves a mixture of the
# patterns it lies closest to instead of snapping to one; a single update is
# exactly attention with the stored patterns as both keys and values.
hopfield(64, beta=0.25, steps=1, name="memory")
relu()
# Width inferred from the output contract: 10 classes.
linear(name="logits")
Input ['B', 1, 28, 28] → output ['B', 10].
Network: [B, 1, 28, 28] -> [B, 10] dtype=float32
index name operation input shapes output shapes
0 n0 flatten x=[B, 1, 28, 28] out=[B, 784]
1 query linear x=[B, 784] out=[B, 128]
2 memory hopfield x=[B, 128] out=[B, 128]
3 n3 relu x=[B, 128] out=[B, 128]
4 logits linear x=[B, 128] out=[B, 10]
Parameters: 109,962. Hopfield layer of Ramsauer et al., "Hopfield Networks is All You Need" (2020), whose only parameter is the stored-pattern matrix; 784·128+128 + 64·128 + 128·10+10 = 109,962 parameters.
LeNet-5¶
The classic convolutional digit classifier: two 5×5 convolution and 2×2 average-pooling stages, then 120- and 84-unit tanh layers. The final width is inferred from the output contract.
examples/networks/lenet5.hndl
# LeNet-5 (LeCun et al. 1998) for 32×32 grayscale digits.
# Valid 5×5 convolutions and 2×2 average pooling: 32 -> 28 -> 14 -> 10 -> 5.
# C1/S2: six feature maps, then subsampling.
conv(6, kernel_size=5, name="c1")
tanh()
avg_pool(2, name="s2")
# C3/S4: sixteen feature maps, then subsampling. Leaves [B, 16, 5, 5].
conv(16, kernel_size=5, name="c3")
tanh()
avg_pool(2, name="s4")
# C5/F6: the classifier head over the 400 flattened features.
flatten()
linear(120, name="c5")
tanh()
linear(84, name="f6")
tanh()
# The output width is inferred from the contract (10 classes).
linear(name="output")
Input ['B', 1, 32, 32] → output ['B', 10].
Network: [B, 1, 32, 32] -> [B, 10] dtype=float32
index name operation input shapes output shapes
0 c1 conv x=[B, 1, 32, 32] out=[B, 6, 28, 28]
1 n1 tanh x=[B, 6, 28, 28] out=[B, 6, 28, 28]
2 s2 avg_pool x=[B, 6, 28, 28] out=[B, 6, 14, 14]
3 c3 conv x=[B, 6, 14, 14] out=[B, 16, 10, 10]
4 n4 tanh x=[B, 16, 10, 10] out=[B, 16, 10, 10]
5 s4 avg_pool x=[B, 16, 10, 10] out=[B, 16, 5, 5]
6 n6 flatten x=[B, 16, 5, 5] out=[B, 400]
7 c5 linear x=[B, 400] out=[B, 120]
8 n8 tanh x=[B, 120] out=[B, 120]
9 f6 linear x=[B, 120] out=[B, 84]
10 n10 tanh x=[B, 84] out=[B, 84]
11 output linear x=[B, 84] out=[B, 10]
Parameters: 61,706. LeCun et al., "Gradient-Based Learning Applied to Document Recognition" (Proc. IEEE, 1998), in the usual modern reading with fully connected C5/F6 layers; checked against the equivalent torch.nn stack: 156 + 2,416 + 48,120 + 10,164 + 850 = 61,706 parameters.
Multilayer perceptron¶
Flatten a 28×28 image and classify it with two ReLU hidden layers. The final width is inferred from the output contract.
examples/networks/mlp.hndl
# A two-hidden-layer perceptron for 28×28 grayscale images (MNIST-sized).
flatten()
linear(256)
relu()
dropout(0.2)
linear(128)
relu()
linear()
Input ['B', 1, 28, 28] → output ['B', 10].
Network: [B, 1, 28, 28] -> [B, 10] dtype=float32
index name operation input shapes output shapes
0 n0 flatten x=[B, 1, 28, 28] out=[B, 784]
1 n1 linear x=[B, 784] out=[B, 256]
2 n2 relu x=[B, 256] out=[B, 256]
3 n3 dropout x=[B, 256] out=[B, 256]
4 n4 linear x=[B, 256] out=[B, 128]
5 n5 relu x=[B, 128] out=[B, 128]
6 n6 linear x=[B, 128] out=[B, 10]
Parameters: 235,146. Classic MNIST baseline; 784·256+256 + 256·128+128 + 128·10+10 = 235,146 parameters.
Mixture-of-experts transformer¶
A causal language model on 32 token ids: learned token and position embeddings, two pre-norm blocks that pair 4-head causal attention with a top-2-of-4 sparse mixture of experts, then a final norm and a bias-free vocabulary head whose width comes from the output contract.
examples/networks/moe_transformer.hndl
# A causal language model whose feed-forward is a sparse mixture of experts.
# Two pre-norm blocks over 32 token positions of a 100-word vocabulary.
# Stem: token ids -> 64-wide sequence, plus a learned absolute position table.
embedding(100, 64, name="tokens")
x0 = pos_embed(32, name="positions")
# Block 1. Pre-norm: normalize a copy, then add the unnormalized stream back.
h = layer_norm(x0, name="block1_attn_norm")
a = attention(h, 4, causal=True, name="block1_attn") # 4 heads of width 16
r = add(a, x0, name="block1_attn_residual")
h2 = layer_norm(r, name="block1_moe_norm")
m = moe(h2, 4, 128, top_k=2, name="block1_moe") # 4 experts, 2 per token
x1 = add(m, r, name="block1_moe_residual")
# Block 2, identical in shape; every layer is its own module.
h3 = layer_norm(x1, name="block2_attn_norm")
a2 = attention(h3, 4, causal=True, name="block2_attn")
r2 = add(a2, x1, name="block2_attn_residual")
h4 = layer_norm(r2, name="block2_moe_norm")
m2 = moe(h4, 4, 128, top_k=2, name="block2_moe")
x2 = add(m2, r2, name="block2_moe_residual")
# Final norm and the untied output head; its width is the vocabulary size,
# which the output contract already fixes at 100.
layer_norm(x2, name="final_norm")
linear(bias=False, name="lm_head")
Input ['B', 32] (input_dtype="int64") → output ['B', 32, 100].
Network: [B, 32] -> [B, 32, 100] dtype=float32 input_dtype=int64
index name operation input shapes output shapes
0 tokens embedding ids=[B, 32] out=[B, 32, 64]
1 positions pos_embed x=[B, 32, 64] out=[B, 32, 64]
2 block1_attn_norm layer_norm x=[B, 32, 64] out=[B, 32, 64]
3 block1_attn attention x=[B, 32, 64] out=[B, 32, 64]
4 block1_attn_residual add a=[B, 32, 64], b=[B, 32, 64] out=[B, 32, 64]
5 block1_moe_norm layer_norm x=[B, 32, 64] out=[B, 32, 64]
6 block1_moe moe x=[B, 32, 64] out=[B, 32, 64]
7 block1_moe_residual add a=[B, 32, 64], b=[B, 32, 64] out=[B, 32, 64]
8 block2_attn_norm layer_norm x=[B, 32, 64] out=[B, 32, 64]
9 block2_attn attention x=[B, 32, 64] out=[B, 32, 64]
10 block2_attn_residual add a=[B, 32, 64], b=[B, 32, 64] out=[B, 32, 64]
11 block2_moe_norm layer_norm x=[B, 32, 64] out=[B, 32, 64]
12 block2_moe moe x=[B, 32, 64] out=[B, 32, 64]
13 block2_moe_residual add a=[B, 32, 64], b=[B, 32, 64] out=[B, 32, 64]
14 final_norm layer_norm x=[B, 32, 64] out=[B, 32, 64]
15 lm_head linear x=[B, 32, 64] out=[B, 32, 100]
Parameters: 181,888. Top-k expert routing follows Shazeer et al. 2017 ("Outrageously Large Neural Networks") and the Switch Transformer (Fedus et al. 2021); the count was checked against an equivalent plain-PyTorch build: embedding 100·64 = 6,400, positions 32·64 = 2,048, five layer norms 5·128 = 640, two attentions 2·4·(64·64+64) = 33,280, two MoE layers 2·(4·64 + 4·(128·64+128 + 64·128+64)) = 133,120, head 100·64 = 6,400, total 181,888.
ResNet-18¶
The 18-layer residual network: a 7×7/2 stem with max pooling, then four stages of two basic blocks at 64, 128, 256 and 512 channels, a global average pool, and one classifier whose 1000 outputs come from the output contract.
examples/networks/resnet18.hndl
# ResNet-18 (He et al., 2015) for 224x224 ImageNet classification.
# Four stages of two basic blocks each; every stage but the first halves
# height and width and doubles the channel count, ending at 512 on 7x7.
# Stem: 7x7/2 convolution, then 3x3/2 max pooling. 224 -> 112 -> 56.
# The convolution omits its bias because the batch norm that follows cancels it.
conv(64, kernel_size=7, stride=2, padding=3, bias=False, name="stem_conv")
batch_norm(name="stem_norm")
relu()
max_pool(3, stride=2, padding=1, name="stem_pool")
# layer1: 64 channels at 56x56. Both shortcuts are identities - stride 1, same width.
resblock(64, name="layer1_0")
resblock(64, name="layer1_1")
# layer2: 128 channels at 28x28. A strided block projects its shortcut with 1x1 conv + norm.
resblock(128, stride=2, name="layer2_0")
resblock(128, name="layer2_1")
# layer3: 256 channels at 14x14.
resblock(256, stride=2, name="layer3_0")
resblock(256, name="layer3_1")
# layer4: 512 channels at 7x7.
resblock(512, stride=2, name="layer4_0")
resblock(512, name="layer4_1")
# Head: average each of the 512 channels over 7x7, then one classifier.
# Its 1000 outputs are inferred from the output contract.
global_avg_pool()
linear(name="fc")
Input ['B', 3, 224, 224] → output ['B', 1000].
Network: [B, 3, 224, 224] -> [B, 1000] dtype=float32
index name operation input shapes output shapes
0 stem_conv conv x=[B, 3, 224, 224] out=[B, 64, 112, 112]
1 stem_norm batch_norm x=[B, 64, 112, 112] out=[B, 64, 112, 112]
2 n2 relu x=[B, 64, 112, 112] out=[B, 64, 112, 112]
3 stem_pool max_pool x=[B, 64, 112, 112] out=[B, 64, 56, 56]
4 layer1_0 resblock x=[B, 64, 56, 56] out=[B, 64, 56, 56]
5 layer1_1 resblock x=[B, 64, 56, 56] out=[B, 64, 56, 56]
6 layer2_0 resblock x=[B, 64, 56, 56] out=[B, 128, 28, 28]
7 layer2_1 resblock x=[B, 128, 28, 28] out=[B, 128, 28, 28]
8 layer3_0 resblock x=[B, 128, 28, 28] out=[B, 256, 14, 14]
9 layer3_1 resblock x=[B, 256, 14, 14] out=[B, 256, 14, 14]
10 layer4_0 resblock x=[B, 256, 14, 14] out=[B, 512, 7, 7]
11 layer4_1 resblock x=[B, 512, 7, 7] out=[B, 512, 7, 7]
12 n12 global_avg_pool x=[B, 512, 7, 7] out=[B, 512]
13 fc linear x=[B, 512] out=[B, 1000]
Parameters: 11,689,512. He et al., "Deep Residual Learning for Image Recognition" (CVPR 2016), configuration ResNet-18; the count equals torchvision.models.resnet18(), sum(p.numel() for p in m.parameters()) == 11,689,512, because hndl's resblock is torchvision's BasicBlock — bias-free 3×3 convolutions with batch norm, and a 1×1 convolution plus norm on each downsampling shortcut.
SAGAN self-attention block¶
The self-attention block of the Self-Attention GAN over a 64-channel 16×16 feature map, written as the single spatial_attention operator: three 1×1 projections, a 256×256 attention map, and a learned gate that starts at zero so the block begins as the identity. The default reduction of 8 gives the paper's C/8 query width, and the batch stays symbolic throughout.
examples/networks/sagan_attention.hndl
# A SAGAN self-attention block (Zhang et al. 2018, section 3) over a
# [B, 64, 16, 16] feature map. Every position attends to every other, so the
# block sees the whole map where a 3x3 convolution sees a neighbourhood.
#
# N = 16*16 = 256 positions, C = 64 channels, and the query/key width is
# C/8 = 8 as in the paper: reduction=8 is the default.
#
# The one operator carries the whole block — the three 1x1 projections, the
# 256x256 attention map, and the learned gate that starts at 0 so the block
# begins as the identity and the network decides how much attention to admit.
spatial_attention()
Input ['B', 64, 16, 16] → output ['B', 64, 16, 16].
Network: [B, 64, 16, 16] -> [B, 64, 16, 16] dtype=float32
index name operation input shapes output shapes
0 n0 spatial_attention x=[B, 64, 16, 16] out=[B, 64, 16, 16]
Parameters: 5,201. Zhang, Goodfellow, Metaxas & Odena, "Self-Attention Generative Adversarial Networks" (ICML 2019, arXiv 2018), section 3: f and g project to C/8 = 8 channels, h keeps C = 64, the attention map is softmax(f(x)^T g(x)) over the key positions, and the output is y = gamma * o + x with gamma initialized to 0. The count matches the same block in plain torch — three Conv2d(1×1) layers with bias (64·8+8 twice and 64·64+64) plus the single gamma — which also gives 5,201.
Small U-Net¶
An encoder–decoder with skip connections for 32×32 segmentation: two 16- and 32-channel down stages, a 64-channel bottleneck, and two up stages that concatenate the matching encoder tensor before convolving. The single output channel of the 1×1 head is inferred from the output contract.
examples/networks/unet_small.hndl
# A small U-Net for 32×32 segmentation: two down stages, a bottleneck, two up
# stages. Every 3×3 convolution uses padding 1, so only the pools and the
# upsamples change the resolution.
# --- Encoder stage 1: 32×32, 16 channels --------------------------------------
conv(16, kernel_size=3, padding=1)
relu()
conv(16, kernel_size=3, padding=1)
skip1 = relu(name="skip1") # [B, 16, 32, 32], kept for the last merge
max_pool(2) # -> 16×16
# --- Encoder stage 2: 16×16, 32 channels --------------------------------------
conv(32, kernel_size=3, padding=1)
relu()
conv(32, kernel_size=3, padding=1)
skip2 = relu(name="skip2") # [B, 32, 16, 16], kept for the first merge
max_pool(2) # -> 8×8
# --- Bottleneck: 8×8, 64 channels ---------------------------------------------
conv(64, kernel_size=3, padding=1)
relu()
conv(64, kernel_size=3, padding=1)
relu()
# --- Decoder stage 2: back to 16×16 -------------------------------------------
upsample(2) # nearest-neighbour, no parameters
conv(32, kernel_size=3, padding=1) # the "up-conv" that halves the channels
upconv2 = relu(name="upconv2")
concat(skip2, upconv2) # 32 + 32 -> 64 channels
conv(32, kernel_size=3, padding=1)
relu()
conv(32, kernel_size=3, padding=1)
relu()
# --- Decoder stage 1: back to 32×32 -------------------------------------------
upsample(2)
conv(16, kernel_size=3, padding=1)
upconv1 = relu(name="upconv1")
concat(skip1, upconv1) # 16 + 16 -> 32 channels
conv(16, kernel_size=3, padding=1)
relu()
conv(16, kernel_size=3, padding=1)
relu()
# --- Head: one logit per pixel; the single channel comes from the contract ----
conv(kernel_size=1, name="logits")
Input ['B', 3, 32, 32] → output ['B', 1, 32, 32].
Network: [B, 3, 32, 32] -> [B, 1, 32, 32] dtype=float32
index name operation input shapes output shapes
0 n0 conv x=[B, 3, 32, 32] out=[B, 16, 32, 32]
1 n1 relu x=[B, 16, 32, 32] out=[B, 16, 32, 32]
2 n2 conv x=[B, 16, 32, 32] out=[B, 16, 32, 32]
3 skip1 relu x=[B, 16, 32, 32] out=[B, 16, 32, 32]
4 n4 max_pool x=[B, 16, 32, 32] out=[B, 16, 16, 16]
5 n5 conv x=[B, 16, 16, 16] out=[B, 32, 16, 16]
6 n6 relu x=[B, 32, 16, 16] out=[B, 32, 16, 16]
7 n7 conv x=[B, 32, 16, 16] out=[B, 32, 16, 16]
8 skip2 relu x=[B, 32, 16, 16] out=[B, 32, 16, 16]
9 n9 max_pool x=[B, 32, 16, 16] out=[B, 32, 8, 8]
10 n10 conv x=[B, 32, 8, 8] out=[B, 64, 8, 8]
11 n11 relu x=[B, 64, 8, 8] out=[B, 64, 8, 8]
12 n12 conv x=[B, 64, 8, 8] out=[B, 64, 8, 8]
13 n13 relu x=[B, 64, 8, 8] out=[B, 64, 8, 8]
14 n14 upsample x=[B, 64, 8, 8] out=[B, 64, 16, 16]
15 n15 conv x=[B, 64, 16, 16] out=[B, 32, 16, 16]
16 upconv2 relu x=[B, 32, 16, 16] out=[B, 32, 16, 16]
17 n17 concat x0=[B, 32, 16, 16], x1=[B, 32, 16, 16] out=[B, 64, 16, 16]
18 n18 conv x=[B, 64, 16, 16] out=[B, 32, 16, 16]
19 n19 relu x=[B, 32, 16, 16] out=[B, 32, 16, 16]
20 n20 conv x=[B, 32, 16, 16] out=[B, 32, 16, 16]
21 n21 relu x=[B, 32, 16, 16] out=[B, 32, 16, 16]
22 n22 upsample x=[B, 32, 16, 16] out=[B, 32, 32, 32]
23 n23 conv x=[B, 32, 32, 32] out=[B, 16, 32, 32]
24 upconv1 relu x=[B, 16, 32, 32] out=[B, 16, 32, 32]
25 n25 concat x0=[B, 16, 32, 32], x1=[B, 16, 32, 32] out=[B, 32, 32, 32]
26 n26 conv x=[B, 32, 32, 32] out=[B, 16, 32, 32]
27 n27 relu x=[B, 16, 32, 32] out=[B, 16, 32, 32]
28 n28 conv x=[B, 16, 32, 32] out=[B, 16, 32, 32]
29 n29 relu x=[B, 16, 32, 32] out=[B, 16, 32, 32]
30 logits conv x=[B, 16, 32, 32] out=[B, 1, 32, 32]
Parameters: 129,841. Ronneberger, Fischer & Brox, "U-Net: Convolutional Networks for Biomedical Image Segmentation" (MICCAI 2015), scaled down to 16/32/64 channels with padded 3×3 convolutions and nearest-neighbour up-convolutions; the count was checked against the same architecture written as a plain torch.nn.Module (448 + 2320 + 4640 + 9248 + 18496 + 36928 + 18464 + 18464 + 9248 + 4624 + 4624 + 2320 + 17 = 129,841 over its 13 convolutions).
Tiny Vision Transformer¶
A ViT-Tiny-shaped classifier for 32×32 RGB images: a 4×4 patch stem of width 192, a learned class token and position table, four pre-norm transformer blocks with 3 heads, then a final layer norm and a linear head read off the class position.
examples/networks/vit_tiny.hndl
# A tiny Vision Transformer (ViT) for 32x32 RGB images.
# Width 192, depth 4, 3 heads (head width 64) — ViT-Tiny proportions at CIFAR scale.
# Stem: 4x4 patches of a 32x32 image give (32/4)*(32/4) = 64 tokens of width 192.
patch_embed(192, 4, name="patches")
# Prepend the learned [CLS] summary position, then add absolute positions.
cls_token(name="cls")
pos_embed(65) # 64 patch tokens + the class token
# Encoder: four pre-norm blocks, gelu feed-forward of width 4*192 = 768.
transformer_block(3, mlp_ratio=4)
transformer_block(3, mlp_ratio=4)
transformer_block(3, mlp_ratio=4)
transformer_block(3, mlp_ratio=4)
# Head: normalize, read the class position, classify.
layer_norm(name="final_norm")
pool_tokens("first")
linear() # 10 classes, inferred from the output contract
Input ['B', 3, 32, 32] → output ['B', 10].
Network: [B, 3, 32, 32] -> [B, 10] dtype=float32
index name operation input shapes output shapes
0 patches patch_embed x=[B, 3, 32, 32] out=[B, 64, 192]
1 cls cls_token x=[B, 64, 192] out=[B, 65, 192]
2 n2 pos_embed x=[B, 65, 192] out=[B, 65, 192]
3 n3 transformer_block x=[B, 65, 192] out=[B, 65, 192]
4 n4 transformer_block x=[B, 65, 192] out=[B, 65, 192]
5 n5 transformer_block x=[B, 65, 192] out=[B, 65, 192]
6 n6 transformer_block x=[B, 65, 192] out=[B, 65, 192]
7 final_norm layer_norm x=[B, 65, 192] out=[B, 65, 192]
8 n8 pool_tokens x=[B, 65, 192] out=[B, 192]
9 n9 linear x=[B, 192] out=[B, 10]
Parameters: 1,803,850. Dosovitskiy et al. 2020, "An Image Is Worth 16x16 Words" (ViT), in the ViT-Tiny configuration (width 192, 3 heads, mlp_ratio 4) at depth 4 with a 4×4 patch stem on 32×32 inputs; checked against an equivalent plain-PyTorch build: 9,408 stem + 192 class token + 12,480 positions + 4·444,864 blocks + 384 final norm + 1,930 head = 1,803,850.