# Memory-footprint knob (set before JAX initialises its GPU backend). At
# 224x224 the VGG/NiN convolutions' true peak is dominated not by activations
# but by the cuDNN convolution workspace XLA allocates while autotuning
# algorithms. Capping XLA's memory pool bounds that workspace budget, so
# autotuning simply picks the fastest algorithm that fits -- cutting the true
# peak from ~7.6 GiB to ~6.4 GiB with numerically identical convolutions and
# no slowdown (unlike disabling autotune outright).
import os
os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '.25'
from d2l import jax as d2l
from flax import nnx
from jax import numpy as jnpBlocks, Bottlenecks, and Branches: VGG, NiN, GoogLeNet
After AlexNet: organize the convolutions
AlexNet proved deep CNNs work, but gave no template: every layer was designed individually.
The next generation contributed one organizing idea each:
- VGG repeats identical blocks; a network becomes a tuple.
- NiN mixes channels with 1×1 convs and replaces the FC head with global average pooling.
- GoogLeNet runs branches of several filter sizes in parallel and concatenates.
All three ideas are still in every modern network.
VGG: regular blocks at scale
VGG (Simonyan & Zisserman, 2014) is AlexNet taken seriously: stack more layers, but make them regular blocks of 3×3 conv + ReLU, ending in a 2×2 max-pool.
From AlexNet’s hand-tuned layers to VGG’s repeated 3×3 blocks.
Receptive field arithmetic
Why 3×3 only? Stacking small kernels grows the visible patch without paying for a large kernel in one step. For stride 1:
r = 1 + \sum_{i=1}^{L} (k_i - 1).
Two 3×3 convolutions see 1 + 2 + 2 = 5 pixels across: the same receptive field as one 5×5 conv, with 18c^2 instead of 25c^2 weights and one extra ReLU. Deep-and-narrow beats shallow-and-wide.
The VGG block
A reusable subunit: num_convs consecutive Conv-ReLU pairs, then a 2×2 MaxPool:
def vgg_block(num_convs, in_channels, out_channels, rngs):
layers = []
for _ in range(num_convs):
layers.append(nnx.Conv(in_channels, out_channels,
kernel_size=(3, 3), padding=(1, 1), rngs=rngs))
layers.append(nnx.relu)
in_channels = out_channels
layers.append(lambda x: nnx.max_pool(
x, window_shape=(2, 2), strides=(2, 2)))
return nnx.Sequential(*layers)The VGG network
Five blocks at growing channel counts plus a 3-layer dense head. The “named architecture” is just a tuple of (n_convs, channels) pairs; a different tuple gives VGG-13/16/19:
class VGG(d2l.Classifier):
def __init__(self, arch, lr=0.1, num_classes=10,
input_shape=(224, 224, 1), rngs=None):
super().__init__()
self.save_hyperparameters(ignore=['rngs'])
rngs = (nnx.Rngs(params=d2l.get_key(), dropout=d2l.get_key())
if rngs is None else rngs)
conv_blks = []
in_channels = input_shape[-1]
for num_convs, out_channels in arch:
conv_blks.append(vgg_block(
num_convs, in_channels, out_channels, rngs))
in_channels = out_channels
height, width = input_shape[:2]
flat_features = (height // 2 ** len(arch)) * (
width // 2 ** len(arch)) * in_channels
self.net = nnx.Sequential(
*conv_blks,
lambda x: x.reshape((x.shape[0], -1)), # flatten
nnx.Linear(flat_features, 4096, rngs=rngs), nnx.relu,
nnx.Dropout(0.5, rngs=rngs),
nnx.Linear(4096, 4096, rngs=rngs), nnx.relu,
nnx.Dropout(0.5, rngs=rngs),
nnx.Linear(4096, num_classes, rngs=rngs))VGG-11 shape check
Each block halves the resolution; channels double until 512:
Sequential output shape: (1, 112, 112, 64)
Sequential output shape: (1, 56, 56, 128)
Sequential output shape: (1, 28, 28, 256)
Sequential output shape: (1, 14, 14, 512)
Sequential output shape: (1, 7, 7, 512)
function output shape: (1, 25088)
...
custom_jvp output shape: (1, 4096)
Dropout output shape: (1, 4096)
Linear output shape: (1, 4096)
custom_jvp output shape: (1, 4096)
Dropout output shape: (1, 4096)
Linear output shape: (1, 10)Training a thin VGG
Full VGG-11 is heavy for a notebook, so we thin the channels (16/32/64/128/128) and train on Fashion-MNIST:
Same pipeline as AlexNet; the block design is what changed.
NiN: 1×1 convolutions and GAP
Network in network (Lin et al., 2013) attacks the FC head: VGG-11’s first dense layer alone needs ~400 MB in FP32.
- 1×1 convolutions: an MLP applied at every pixel; channel mixing and extra nonlinearity at zero spatial cost.
- Global average pooling: one number per channel replaces the giant dense classifier.
NiN vs. VGG: same body idea, radically different head.
The NiN block
A regular convolution followed by two 1×1 convolutions with ReLUs in between:
def nin_block(in_channels, out_channels, kernel_size, strides, padding, rngs):
return nnx.Sequential(
nnx.Conv(in_channels, out_channels, kernel_size, strides=strides,
padding=padding, rngs=rngs), nnx.relu,
nnx.Conv(out_channels, out_channels, kernel_size=(1, 1), rngs=rngs),
nnx.relu,
nnx.Conv(out_channels, out_channels, kernel_size=(1, 1), rngs=rngs),
nnx.relu)The NiN model
Four NiN blocks (AlexNet’s kernel sizes), max-pool between them, and the last block emits num_classes channels. Then global average pooling. No fully connected layers at all.
class NiN(d2l.Classifier):
def __init__(self, lr=0.1, num_classes=10, rngs=None):
super().__init__()
self.save_hyperparameters(ignore=['rngs'])
rngs = (nnx.Rngs(params=d2l.get_key(), dropout=d2l.get_key())
if rngs is None else rngs)
self.net = nnx.Sequential(
nin_block(1, 96, (11, 11), (4, 4), (0, 0), rngs),
lambda x: nnx.max_pool(x, (3, 3), strides=(2, 2)),
nin_block(96, 256, (5, 5), (1, 1), (2, 2), rngs),
lambda x: nnx.max_pool(x, (3, 3), strides=(2, 2)),
nin_block(256, 384, (3, 3), (1, 1), (1, 1), rngs),
lambda x: nnx.max_pool(x, (3, 3), strides=(2, 2)),
nnx.Dropout(0.5, rngs=rngs),
nin_block(384, num_classes, (3, 3), (1, 1), (1, 1), rngs),
lambda x: x.mean(axis=(1, 2)), # global avg pooling over H, W (NHWC)
lambda x: x.reshape((x.shape[0], -1))) # flattenNiN shape inspection
Spatial dims shrink, channels grow, and the final block already has one channel per class:
Sequential output shape: (1, 54, 54, 96)
function output shape: (1, 26, 26, 96)
Sequential output shape: (1, 26, 26, 256)
function output shape: (1, 12, 12, 256)
Sequential output shape: (1, 12, 12, 384)
function output shape: (1, 5, 5, 384)
Dropout output shape: (1, 5, 5, 384)
Sequential output shape: (1, 5, 5, 10)
function output shape: (1, 10)
function output shape: (1, 10)Training NiN
The averaging head costs nothing and does not hurt accuracy; that surprise made GAP the default head ever since.
GoogLeNet: go wide
GoogLeNet (Szegedy et al., 2015) won ImageNet 2014 with two lasting contributions:
- The stem / body / head decomposition of a network, still the universal vocabulary.
- The Inception block: don’t pick a filter size, run 1×1, 3×3, 5×5, and pooling in parallel and concatenate the channels.
Four branches, four scales, one shared input.
The Inception block in code
Four branches, channel-concatenated; 1×1 convs shrink channels before the costly 3×3 and 5×5:
class Inception(nnx.Module):
def __init__(self, in_channels, c1, c2, c3, c4, rngs):
# Branch 1
self.b1_1 = nnx.Conv(in_channels, c1, kernel_size=(1, 1), rngs=rngs)
# Branch 2
self.b2_1 = nnx.Conv(in_channels, c2[0], kernel_size=(1, 1), rngs=rngs)
self.b2_2 = nnx.Conv(c2[0], c2[1], kernel_size=(3, 3),
padding='same', rngs=rngs)
# Branch 3
self.b3_1 = nnx.Conv(in_channels, c3[0], kernel_size=(1, 1), rngs=rngs)
self.b3_2 = nnx.Conv(c3[0], c3[1], kernel_size=(5, 5),
padding='same', rngs=rngs)
# Branch 4
self.b4_2 = nnx.Conv(in_channels, c4, kernel_size=(1, 1), rngs=rngs)
def __call__(self, x):
b1 = nnx.relu(self.b1_1(x))
b2 = nnx.relu(self.b2_2(nnx.relu(self.b2_1(x))))
b3 = nnx.relu(self.b3_2(nnx.relu(self.b3_1(x))))
pooled = nnx.max_pool(x, window_shape=(3, 3),
strides=(1, 1), padding='same')
b4 = nnx.relu(self.b4_2(pooled))
return jnp.concatenate((b1, b2, b3, b4), axis=-1)Bottleneck arithmetic
First body block: 192 channels in, 64+128+32+32 = 256 out, spatial size unchanged.
The 5×5 branch, direct: 25 \cdot 192 \cdot 32 \approx 154k weights.
With a 16-channel 1×1 bottleneck first: 192 \cdot 16 + 25 \cdot 16 \cdot 32 \approx 16k, a 10× saving. This trick is in nearly every network since.
The whole network as data
Nine Inception blocks in three groups (2, 5, 2), pooling between groups. The hand-picked channel allocations are just a tuple; assembly is stem + body + head:
arch = (((64, (96, 128), (16, 32), 32), (128, (128, 192), (32, 96), 64)),
((192, (96, 208), (16, 48), 64), (160, (112, 224), (24, 64), 64),
(128, (128, 256), (24, 64), 64), (112, (144, 288), (32, 64), 64),
(256, (160, 320), (32, 128), 128)),
((256, (160, 320), (32, 128), 128), (384, (192, 384), (48, 128), 128)))class GoogleNet(d2l.Classifier):
def __init__(self, lr=0.1, num_classes=10, rngs=None):
super().__init__()
self.save_hyperparameters(ignore=['rngs'])
rngs = nnx.Rngs(d2l.get_key()) if rngs is None else rngs
pool = lambda x: nnx.max_pool(
x, window_shape=(3, 3), strides=(2, 2), padding='same')
stem = nnx.Sequential(
nnx.Conv(1, 64, kernel_size=(7, 7), strides=(2, 2),
padding='same', rngs=rngs), nnx.relu, pool,
nnx.Conv(64, 64, kernel_size=(1, 1), rngs=rngs), nnx.relu,
nnx.Conv(64, 192, kernel_size=(3, 3), padding='same', rngs=rngs),
nnx.relu, pool)
body, in_channels = [], 192
for group in arch:
blocks = []
for c1, c2, c3, c4 in group:
blocks.append(Inception(
in_channels, c1, c2, c3, c4, rngs))
in_channels = c1 + c2[1] + c3[1] + c4
body.append(nnx.Sequential(*blocks))
head = nnx.Sequential(lambda x: x.mean(axis=(1, 2)),
nnx.Linear(in_channels, num_classes, rngs=rngs))
self.net = nnx.Sequential(stem, body[0], pool, body[1], pool,
body[2], head)Shape check
Sequential output shape: (1, 12, 12, 192)
Sequential output shape: (1, 12, 12, 480)
function output shape: (1, 6, 6, 480)
Sequential output shape: (1, 6, 6, 832)
function output shape: (1, 3, 3, 832)
Sequential output shape: (1, 3, 3, 1024)
Sequential output shape: (1, 10)Cheaper than VGG (~7M vs. ~138M parameters) and more accurate: the start of deliberate cost–accuracy design.
What survived
- Blocks (VGG): everything is specified block by block.
- 1×1 convolution (NiN): the standard channel mixer.
- Global average pooling (NiN): the default head.
- Multi-branch (GoogLeNet): survives as grouped convolutions (ResNeXt) and train-time-only branches (RepVGG).
- Inception-style hand-tuned branch cocktails: extinct.
Next ingredient: normalization.