from d2l import jax as d2l
from flax import nnx
from jax import numpy as jnpClassical vision pipelines never fed raw pixels to a classifier:
Progress meant inventing better features, not better learning.
LeCun, Hinton, Bengio, Ng, Amari, Schmidhuber: features should be learned, hierarchically, layer by layer.
AlexNet’s first layer learned filters that resemble the hand-crafted ones:
First-layer filters learned by AlexNet.
AlexNet (Krizhevsky, Sutskever, Hinton, 2012) put them together and won ILSVRC 2012 by a large margin.
Same design, scaled up: convolutional stages, then a fully connected head.
LeNet and AlexNet side by side.
Five conv layers (11×11 → 5×5 → three 3×3) with max-pooling, then two 4096-wide dense layers with dropout:
class AlexNet(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(
nnx.Conv(1, 96, kernel_size=(11, 11), strides=4,
padding='VALID', rngs=rngs),
nnx.relu,
lambda x: nnx.max_pool(x, window_shape=(3, 3), strides=(2, 2)),
nnx.Conv(96, 256, kernel_size=(5, 5), rngs=rngs),
nnx.relu,
lambda x: nnx.max_pool(x, window_shape=(3, 3), strides=(2, 2)),
nnx.Conv(256, 384, kernel_size=(3, 3), rngs=rngs), nnx.relu,
nnx.Conv(384, 384, kernel_size=(3, 3), rngs=rngs), nnx.relu,
nnx.Conv(384, 256, kernel_size=(3, 3), rngs=rngs), nnx.relu,
lambda x: nnx.max_pool(x, window_shape=(3, 3), strides=(2, 2)),
lambda x: x.reshape((x.shape[0], -1)), # flatten
nnx.Linear(5 * 5 * 256, 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))Walk a single 224×224 image through the network and print each block’s output shape, from 224×224 down to 6×6 at 256 channels:
Conv output shape: (1, 54, 54, 96)
custom_jvp output shape: (1, 54, 54, 96)
function output shape: (1, 26, 26, 96)
Conv output shape: (1, 26, 26, 256)
custom_jvp output shape: (1, 26, 26, 256)
function output shape: (1, 12, 12, 256)
...
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)
Upsample the 28×28 Fashion-MNIST images to the 224×224 input AlexNet expects, then train with a smaller learning rate than LeNet: