7.1  The ImageNet Moment: AlexNet

Although CNNs were well known in computer vision and machine learning following the introduction of LeNet (LeCun et al. 1995), they did not immediately dominate the field. For much of the time between the early 1990s and the watershed results of 2012 (Krizhevsky et al. 2012), neural networks were often outperformed by other methods, such as kernel methods (Schölkopf and Smola 2002), ensemble methods (Freund and Schapire 1996), and structured estimation (Taskar et al. 2004).

For computer vision, this comparison is not entirely fair, because practitioners never fed raw pixels into traditional models. A typical pipeline preprocessed the images with hand-crafted feature extractors such as SIFT (the scale-invariant feature transform) (Lowe 2004), SURF (speeded up robust features) (Bay et al. 2006), or bags of visual words (Sivic and Zisserman 2003), and only then trained a linear model or kernel method on the result. The features were crafted rather than learned: progress came from cleverer features and from deep insight into geometry (Hartley and Zisserman 2000), and researchers believed, justifiably, that a slightly bigger or cleaner dataset or a slightly improved feature-extraction pipeline mattered far more to the final accuracy than the choice of classifier. The learning algorithm was an afterthought.

Neural networks also faced real obstacles. The accelerators of the 1990s could not power deep multichannel, multilayer CNNs with many parameters, and datasets were small: OCR on 60,000 low-resolution \(28 \times 28\) pixel images was considered a highly challenging task. Key tricks for training neural networks were still missing, too, including parameter initialization heuristics (Glorot and Bengio 2010), clever variants of stochastic gradient descent (Kingma and Ba 2015), non-squashing activation functions (Nair and Hinton 2010), and effective regularization (Srivastava et al. 2014).

from d2l import torch as d2l
import torch
from torch import nn
from d2l import tensorflow as d2l
import tensorflow as tf
from d2l import jax as d2l
from flax import nnx
from jax import numpy as jnp
from d2l import mxnet as d2l
from mxnet import np, init, npx
from mxnet.gluon import nn
npx.set_np()

7.1.1 Representation Learning

Put differently, the most important part of the classical pipeline was the representation, and up until 2012 it was calculated mostly mechanically: SIFT, SURF, HOG (histograms of oriented gradient) (Dalal and Triggs 2005), and bags of visual words ruled the roost.

Another group of researchers, including Yann LeCun, Geoff Hinton, Yoshua Bengio, Andrew Ng, Shun-ichi Amari, and Juergen Schmidhuber, believed instead that features themselves ought to be learned, hierarchically composed from multiple jointly learned layers. The automatic design of visual features, such as those obtained by sparse coding (Olshausen and Field 1996), remained an open challenge until Dean et al. (2012); Le (2013) and the advent of modern CNNs.

The first modern CNN (Krizhevsky et al. 2012), named AlexNet after one of its inventors, Alex Krizhevsky, is largely an evolutionary improvement over LeNet. It won the 2012 ImageNet challenge and vindicated the bet on learning: in its lowest layers, the network learned feature extractors that resembled the traditional filters, as Figure 7.1.1 shows. Higher layers build upon these representations to capture larger structures, like eyes, noses, and blades of grass, and yet higher layers whole objects, like people, airplanes, and dogs. Ultimately, the final hidden state is a compact representation of the image in which the different categories are easily separated.

Figure 7.1.1: Image filters learned by the first layer of AlexNet. Reproduction courtesy of Krizhevsky et al. (2012).

AlexNet (2012) and its precursor LeNet (1995) share many architectural elements, which raises the question of why it took so long. The decisive difference is that, over the intervening two decades, the amount of data and the computing power available had each grown by orders of magnitude.

7.1.1.1 Missing Ingredient: Data

Deep models with many layers require large amounts of data in order to enter the regime where they significantly outperform traditional methods based on convex optimization. However, given the limited storage, the expense of imaging sensors, and the tighter research budgets of the 1990s, most research relied on tiny datasets of hundreds or a few thousand low-resolution images, such as those in the UCI collection.

The ImageNet dataset, released in 2009 (Deng et al. 2009), changed this. The 2012 classification challenge supplied roughly 1.2 million training images across 1000 categories drawn from WordNet (Miller 1995), prefiltered by web image search and verified by Amazon Mechanical Turk workers. Class sizes varied, and the source images had varying resolutions; models commonly trained on \(224 \times 224\) crops. The scale exceeded earlier labeled datasets by over an order of magnitude, while the source images retained far more detail than the \(32 \times 32\) thumbnails of the 80-million-image TinyImages dataset (Torralba et al. 2008), which allowed higher-level features to form. The associated competition, the ImageNet Large Scale Visual Recognition Challenge (Russakovsky et al. 2015), pushed computer vision and machine learning research to a scale that academics had not previously considered.

7.1.1.2 Missing Ingredient: Hardware

Deep learning models are also voracious consumers of compute cycles: training can take hundreds of epochs, and each iteration passes data through many layers of expensive linear algebra operations. Graphics processing units (GPUs) changed the economics. These chips had been developed to accelerate computer graphics, in particular high-throughput \(4 \times 4\) matrix–vector products, math very similar to that of convolutional layers, and around that time NVIDIA and ATI had begun optimizing them for general computing (Fernando 2004), marketing them as general-purpose GPUs (GPGPUs). Where a CPU core runs at a high clock frequency and spends most of its chip area on the machinery of general control flow (branch predictors, deep pipelines, speculative execution, large caches), a GPU packs thousands of much simpler cores onto one chip. This wins on power, since consumption grows roughly quadratically with clock frequency: for the budget of one CPU core running at four times the speed, 16 GPU cores at \(\frac{1}{4}\) the speed deliver \(16 \times \frac{1}{4} = 4\) times the throughput. GPUs also have far wider memory buses, which matters because many deep learning operations are limited by memory bandwidth. A convolution applies the same small program at many output locations and channels, providing exactly this kind of independent work. The effect compounded quickly: between 1999, when NVIDIA’s GeForce 256 processed roughly 480 million floating-point operations per second with no programming framework beyond graphics APIs, and 2012, consumer GPU throughput grew by roughly three orders of magnitude, and general-purpose GPU interfaces made it accessible without expressing the computation as a graphics pipeline.

This was the situation in 2012 when Alex Krizhevsky and Ilya Sutskever implemented a deep CNN that could run on GPUs. They realized that the computational bottlenecks in CNNs, convolutions and matrix multiplications, are precisely the operations that GPUs parallelize well. Using two NVIDIA GTX 580s with 3 GB of memory each, they implemented fast convolutions. The cuda-convnet code was good enough that for several years it was the industry standard and powered the first couple of years of the deep learning boom.

7.1.2 AlexNet

AlexNet, which employed an 8-layer CNN, won the ImageNet Large Scale Visual Recognition Challenge 2012 by a large margin (Russakovsky et al. 2013). This network showed, for the first time, that the features obtained by learning can transcend manually designed features, breaking the previous paradigm in computer vision.

The architectures of AlexNet and LeNet are closely related, as Figure 7.1.2 illustrates. Note that we provide a slightly streamlined version of AlexNet removing some of the design quirks that were needed in 2012 to make the model fit on two small GPUs.

Figure 7.1.2: LeNet (left) and AlexNet (right) side by side. Both consist of convolutional stages followed by a fully connected head; AlexNet is deeper and wider and replaces sigmoid activations with ReLU.

There are also key differences. First, AlexNet is much deeper than the comparatively small LeNet-5. AlexNet consists of eight layers: five convolutional layers, two fully connected hidden layers, and one fully connected output layer. Second, AlexNet used the ReLU instead of the sigmoid as its activation function. Let’s look at the details.

7.1.2.1 Architecture

In AlexNet’s first layer, the convolution window shape is \(11\times11\). Since the images in ImageNet are eight times taller and wider than the MNIST images, objects in ImageNet data tend to occupy more pixels with more visual detail. Consequently, a larger convolution window is needed to capture the object. The convolution window shape in the second layer is reduced to \(5\times5\), followed by \(3\times3\). In addition, after the first, second, and fifth convolutional layers, the network adds max-pooling layers with a window shape of \(3\times3\) and a stride of 2. Moreover, AlexNet has ten times more convolution channels than LeNet.

After the final convolutional layer, there are two huge fully connected layers with 4096 outputs each. Together they account for almost all of the model’s parameters (over 160 MB in single precision). Because of the limited memory in early GPUs, the original AlexNet used a dual data stream design, so that each of their two GPUs could be responsible for storing and computing only its half of the model. GPU memory is rarely that scarce anymore, so our version of the model dispenses with the split.

7.1.2.2 Activation Functions

AlexNet also changed the sigmoid activation function to the simpler ReLU activation function. This makes the computation cheaper, since the ReLU has no exponentiation operation, and, more importantly, it makes training easier: when the output of the sigmoid is very close to 0 or 1, its gradient is almost 0, so poorly initialized parameters may stop receiving updates, whereas the gradient of the ReLU in the positive interval is always 1 (Section 4.1.2).

7.1.2.3 Capacity Control and Preprocessing

AlexNet controls the model complexity of the fully connected layer by dropout (Section 4.6), while LeNet only uses weight decay. To augment the data even further, the training loop of AlexNet added a great deal of image augmentation, such as flipping, clipping, and color changes. This exposes the model to many more variants of each image and thereby reduces overfitting; see Buslaev et al. (2020) for an in-depth review of such preprocessing steps. Augmentation and its descendants have since grown into a central part of how convolutional networks are trained, and we return to them in Section 7.6.

class AlexNet(d2l.Classifier):
    def __init__(self, lr=0.1, num_classes=10):
        super().__init__()
        self.save_hyperparameters()
        self.net = nn.Sequential(
            nn.LazyConv2d(96, kernel_size=11, stride=4),
            nn.ReLU(), nn.MaxPool2d(kernel_size=3, stride=2),
            nn.LazyConv2d(256, kernel_size=5, padding=2), nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2),
            nn.LazyConv2d(384, kernel_size=3, padding=1), nn.ReLU(),
            nn.LazyConv2d(384, kernel_size=3, padding=1), nn.ReLU(),
            nn.LazyConv2d(256, kernel_size=3, padding=1), nn.ReLU(),
            nn.MaxPool2d(kernel_size=3, stride=2), nn.Flatten(),
            nn.LazyLinear(4096), nn.ReLU(), nn.Dropout(p=0.5),
            nn.LazyLinear(4096), nn.ReLU(),nn.Dropout(p=0.5),
            nn.LazyLinear(num_classes))
        # Note: lazy layers have no parameters at construction time, so weight
        # initialization (d2l.init_cnn) is applied later via apply_init after
        # a dummy forward pass materializes the parameters.
class AlexNet(d2l.Classifier):
    def __init__(self, lr=0.1, num_classes=10):
        super().__init__()
        self.save_hyperparameters()
        self.net = tf.keras.models.Sequential([
            tf.keras.layers.Conv2D(filters=96, kernel_size=11, strides=4,
                                   activation='relu'),
            tf.keras.layers.MaxPool2D(pool_size=3, strides=2),
            tf.keras.layers.Conv2D(filters=256, kernel_size=5, padding='same',
                                   activation='relu'),
            tf.keras.layers.MaxPool2D(pool_size=3, strides=2),
            tf.keras.layers.Conv2D(filters=384, kernel_size=3, padding='same',
                                   activation='relu'),
            tf.keras.layers.Conv2D(filters=384, kernel_size=3, padding='same',
                                   activation='relu'),
            tf.keras.layers.Conv2D(filters=256, kernel_size=3, padding='same',
                                   activation='relu'),
            tf.keras.layers.MaxPool2D(pool_size=3, strides=2),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(4096, activation='relu'),
            tf.keras.layers.Dropout(0.5),
            tf.keras.layers.Dense(4096, activation='relu'),
            tf.keras.layers.Dropout(0.5),
            tf.keras.layers.Dense(num_classes)])
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))
class AlexNet(d2l.Classifier):
    def __init__(self, lr=0.1, num_classes=10):
        super().__init__()
        self.save_hyperparameters()
        self.net = nn.Sequential()
        self.net.add(
            nn.Conv2D(96, kernel_size=11, strides=4, activation='relu'),
            nn.MaxPool2D(pool_size=3, strides=2),
            nn.Conv2D(256, kernel_size=5, padding=2, activation='relu'),
            nn.MaxPool2D(pool_size=3, strides=2),
            nn.Conv2D(384, kernel_size=3, padding=1, activation='relu'),
            nn.Conv2D(384, kernel_size=3, padding=1, activation='relu'),
            nn.Conv2D(256, kernel_size=3, padding=1, activation='relu'),
            nn.MaxPool2D(pool_size=3, strides=2),
            nn.Dense(4096, activation='relu'), nn.Dropout(0.5),
            nn.Dense(4096, activation='relu'), nn.Dropout(0.5),
            nn.Dense(num_classes))
        self.net.initialize(init.Xavier())

We construct a single-channel data example with both height and width of 224 to observe the output shape of each layer. It matches the AlexNet architecture in Figure 7.1.2.

AlexNet().layer_summary((1, 1, 224, 224))
Conv2d output shape:     torch.Size([1, 96, 54, 54])
ReLU output shape:   torch.Size([1, 96, 54, 54])
MaxPool2d output shape:  torch.Size([1, 96, 26, 26])
Conv2d output shape:     torch.Size([1, 256, 26, 26])
ReLU output shape:   torch.Size([1, 256, 26, 26])
MaxPool2d output shape:  torch.Size([1, 256, 12, 12])
Conv2d output shape:     torch.Size([1, 384, 12, 12])
ReLU output shape:   torch.Size([1, 384, 12, 12])
Conv2d output shape:     torch.Size([1, 384, 12, 12])
ReLU output shape:   torch.Size([1, 384, 12, 12])
Conv2d output shape:     torch.Size([1, 256, 12, 12])
ReLU output shape:   torch.Size([1, 256, 12, 12])
MaxPool2d output shape:  torch.Size([1, 256, 5, 5])
Flatten output shape:    torch.Size([1, 6400])
Linear output shape:     torch.Size([1, 4096])
ReLU output shape:   torch.Size([1, 4096])
Dropout output shape:    torch.Size([1, 4096])
Linear output shape:     torch.Size([1, 4096])
ReLU output shape:   torch.Size([1, 4096])
Dropout output shape:    torch.Size([1, 4096])
Linear output shape:     torch.Size([1, 10])
AlexNet().layer_summary((1, 224, 224, 1))
Conv2D output shape:     (1, 54, 54, 96)
MaxPooling2D output shape:   (1, 26, 26, 96)
Conv2D output shape:     (1, 26, 26, 256)
MaxPooling2D output shape:   (1, 12, 12, 256)
Conv2D output shape:     (1, 12, 12, 384)
Conv2D output shape:     (1, 12, 12, 384)
Conv2D output shape:     (1, 12, 12, 256)
MaxPooling2D output shape:   (1, 5, 5, 256)
Flatten output shape:    (1, 6400)
Dense output shape:  (1, 4096)
Dropout output shape:    (1, 4096)
Dense output shape:  (1, 4096)
Dropout output shape:    (1, 4096)
Dense output shape:  (1, 10)
AlexNet().layer_summary((1, 224, 224, 1))
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)
Conv output shape:   (1, 12, 12, 384)
custom_jvp output shape:     (1, 12, 12, 384)
Conv output shape:   (1, 12, 12, 384)
custom_jvp output shape:     (1, 12, 12, 384)
Conv output shape:   (1, 12, 12, 256)
custom_jvp output shape:     (1, 12, 12, 256)
function output shape:   (1, 5, 5, 256)
function output shape:   (1, 6400)
Linear output shape:     (1, 4096)
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)
AlexNet().layer_summary((1, 1, 224, 224))
Conv2D output shape:     (1, 96, 54, 54)
MaxPool2D output shape:  (1, 96, 26, 26)
Conv2D output shape:     (1, 256, 26, 26)
MaxPool2D output shape:  (1, 256, 12, 12)
Conv2D output shape:     (1, 384, 12, 12)
Conv2D output shape:     (1, 384, 12, 12)
Conv2D output shape:     (1, 256, 12, 12)
MaxPool2D output shape:  (1, 256, 5, 5)
Dense output shape:  (1, 4096)
Dropout output shape:    (1, 4096)
Dense output shape:  (1, 4096)
Dropout output shape:    (1, 4096)
Dense output shape:  (1, 10)

7.1.3 Training

Although AlexNet was trained on ImageNet in Krizhevsky et al. (2012), we use Fashion-MNIST here since training an ImageNet model to convergence could take hours or days even on a modern GPU. One of the problems with applying AlexNet directly on Fashion-MNIST is that its images have lower resolution (\(28 \times 28\) pixels) than ImageNet images. To make things work, we upsample them to \(224 \times 224\). This is generally not a smart practice, as it simply increases the computational complexity without adding information. Nonetheless, we do it here to be faithful to the AlexNet architecture. We perform this resizing with the resize argument in the d2l.FashionMNIST constructor.

Now, we can start training AlexNet. Compared to LeNet in Section 6.6, the main change here is the use of a smaller learning rate and much slower training due to the deeper and wider network, the higher image resolution, and the more costly convolutions.

model = AlexNet(lr=0.01)
data = d2l.FashionMNIST(batch_size=128, resize=(224, 224))
trainer = d2l.Trainer(max_epochs=10, num_gpus=1)
if tab.selected('pytorch'):
    # Lazy layers have no weights at construction time; apply_init runs a
    # dummy forward pass to materialize parameters and then applies init_cnn.
    model.apply_init([next(iter(data.get_dataloader(True)))[0]], d2l.init_cnn)
trainer.fit(model, data)

trainer = d2l.Trainer(max_epochs=10)
data = d2l.FashionMNIST(batch_size=128, resize=(224, 224))
with d2l.try_gpu():
    model = AlexNet(lr=0.01)
    trainer.fit(model, data)

model = AlexNet(lr=0.01)
data = d2l.FashionMNIST(batch_size=128, resize=(224, 224))
trainer = d2l.Trainer(max_epochs=10, num_gpus=1)
if tab.selected('pytorch'):
    # Lazy layers have no weights at construction time; apply_init runs a
    # dummy forward pass to materialize parameters and then applies init_cnn.
    model.apply_init([next(iter(data.get_dataloader(True)))[0]], d2l.init_cnn)
trainer.fit(model, data)

model = AlexNet(lr=0.01)
data = d2l.FashionMNIST(batch_size=128, resize=(224, 224))
trainer = d2l.Trainer(max_epochs=10, num_gpus=1)
if tab.selected('pytorch'):
    # Lazy layers have no weights at construction time; apply_init runs a
    # dummy forward pass to materialize parameters and then applies init_cnn.
    model.apply_init([next(iter(data.get_dataloader(True)))[0]], d2l.init_cnn)
trainer.fit(model, data)

7.1.4 Discussion

AlexNet’s structure bears a close resemblance to LeNet, with a number of critical improvements, both for accuracy (dropout) and for ease of training (ReLU). Equally consequential is the progress in deep learning tooling: what was several months of work in 2012 can now be accomplished in a dozen lines of code using any modern framework.

Reviewing the architecture, we see that AlexNet has an Achilles heel when it comes to efficiency: the last two hidden layers require matrices of size \(6400 \times 4096\) and \(4096 \times 4096\), respectively. This corresponds to 164 MB of memory and 81 MFLOPs of computation, both of which are a nontrivial outlay, especially on smaller devices, such as mobile phones. This is one of the reasons why AlexNet has been surpassed by much more effective architectures that we will cover in the following sections. Nonetheless, it is a key step from shallow to deep networks that are used nowadays. Note that even though the number of parameters exceeds by far the amount of training data in our experiments (the last two layers have more than 40 million parameters, trained on a dataset of 60 thousand images), there is hardly any overfitting: training and validation loss are virtually identical throughout training. This is due to the improved regularization, such as dropout, inherent in modern deep network designs.

Although it seems that there are only a few more lines in AlexNet’s implementation than in LeNet’s, it took the academic community many years to embrace this conceptual change and take advantage of its excellent experimental results. This was also due to the lack of efficient computational tools. At the time neither DistBelief (Dean et al. 2012) nor Caffe (Jia et al. 2014) existed, and Theano (Bergstra et al. 2010), the first widely used automatic-differentiation framework, still lacked many features its successors would bring. It was the maturation of such frameworks, from Theano to TensorFlow (Abadi et al. 2016) and later PyTorch (Paszke et al. 2019) and JAX (Frostig et al. 2018), that turned implementing a new architecture from an engineering project into routine work.

7.1.5 Exercises

  1. Following up on the discussion above, analyze the computational properties of AlexNet.
    1. Compute the memory footprint for convolutions and fully connected layers, respectively. Which one dominates?
    2. Calculate the computational cost for the convolutions and the fully connected layers.
    3. How does the memory (read and write bandwidth, latency, size) affect computation? Is there any difference in its effects for training and inference?
  2. You are a chip designer and need to trade off computation and memory bandwidth. For example, a faster chip requires more power and possibly a larger chip area. More memory bandwidth requires more pins and control logic, thus also more area. How do you optimize?
  3. Why do engineers no longer report performance benchmarks on AlexNet?
  4. Try increasing the number of epochs when training AlexNet. Compared with LeNet, how do the results differ? Why?
  5. AlexNet may be too complex for the Fashion-MNIST dataset, in particular due to the low resolution of the initial images.
    1. Try simplifying the model to make the training faster, while ensuring that the accuracy does not drop significantly.
    2. Design a better model that works directly on \(28 \times 28\) images.
  6. Modify the batch size, and observe the changes in throughput (images/s), accuracy, and GPU memory.
  7. Apply dropout and ReLU to LeNet-5. Does it improve? Can you improve things further by preprocessing to take advantage of the invariances inherent in the images?
  8. Can you make AlexNet overfit? Which feature do you need to remove or change to break training?