AlexNet Convolutional Neural Network

Recommended read before jumping in this post:

• Building an Image Classifier from Scratch
• Introduction to Convolutional Neural Network

Our goal for this blog post is to have solid fundamental understanding of the different parts that comprises the AlexNet architecture.

You can see the Colab Notebook with the code implementation here.

AlexNet

It is accurate to say that AlexNet is one of the most important breakthroughs in computer vision, and arguably the catalyst for the modern Deep Learning era.

Pre-AlexNet computer vision relied on manual feature engineering. Human engineers wrote explicit mathematical algorithms to detect specific visual patterns (edges, corners, gradients). These pre-computed features were then fed into a shallow statistical classifier.

AlexNet popularized end-to-end representation learning. Instead of humans deciding what features are important, the model takes raw pixel data and learns the optimal feature-extracting filters itself through training data.

Above is the original AlexNet pipeline. What we will be building at the end is a simplified version of it. But we will be going over the layers (convolution, LRN, max pool, etc) so we will have solid understanding of how they work.

Intuition

As you can see in the architecture, it’s just multiple different layers stacked creating a pipeline. The intuition is our pipeline accepts an image, then as it goes through the convolutional layers, we apply filters (kernels) that scan across the image. These filters learn to detect specific visual patters. Early layers might learn simple features like edges, corners, and textures. Deeper layers combine these simple features to detect more complex shapes and objects like eyes, wheels, mouth, specific textures, etc. The result of each convolution is a set of “feature maps” (height X width X depth/channels) where each point in each channels indicates the presence and strength of a particular feature at that spatial location in the original image. Pooling layers are used to progressively reduce the spatial dimensions (height X width) of the feature maps, making the representation more abstract and computationally efficient while increasing the number of features (depth). After the convolutional / pooling layers, the data is represented as a condensed but very rich, high-dimensional tensor preserving spatial relationships of the detected features.

After several convolutional layers, we flatten the data. Convolutional layers excel at extracting spacial features and preserving their relationship in a multi-dimensional structure. However, the subsequent dense (fully connected layers), which are typically used for classification, operate on 1D vectors. So we need to transition from the 3D tensor structure to a simple, long list of numbers. The flattened data is no longer an image representation in terms of spatial arrangement. Instead, it represent a high-dimensional feature vector. Every single value in this vector corresponds to a specific feature activation (e.g., presence of a specific curve at a specific reduced spatial location, presence of a texture, etc.) across the aggregated feature maps.

Then lastly, the dense (fully connected) layers is where the classification happens. In summary, the convolutional layers extract and condense rich spatial features, flattening translates this into a comprehensive feature vector for classification, and the dense layers learn complex rules and relationships within this feature vector to make final, informed predictions based on all the aggregated visual knowledge.

Now let’s dive deeper to the layers!

Convolutions on RGB images

In our Introduction to Convolutional Neural Network, we focused our example with just a grayscale image (one channel), but we can perform convolutions to more channels. Typically, images have 3 channels (RGB) so our kernel should also have a depth of 3.

To help us understand and visualize the next few sections, let’s consider an input of a 5x5x3 input. This means we have 5 x 5 pixels RGB (x3) image.

(I wrote “1”, “2”, “3” in the kernels to denote that they are different kernels.)

Important: The size of the kernel’s depth should always be the same as the input image’s depth.

The process of performing convolution with more than 1 depth vs a single depth is similar. We place our 3x3x3 kernel exactly over the top-left corner of our 5x5x3 input.

Imagine this is the values of each of the layers in our RGB and the kernel channels. For now, let’s just focus on the first step.

• Row 1: (200 $\times$ 1) + (150 $\times$ 0) + (10 $\times$ -1) = 190
• Row 2: (100 $\times$ 1) + (100 $\times$ 0) + (50 $\times$ -1) = 50
• Row 3: (50 $\times$ 1) + (50 $\times$ 0) + (10 $\times$ -1) = 40
• Red Sum: 190 + 50 + 40 = 280

• Row 1: (25 $\times$ 1) + (25 $\times$ 1) + (25 $\times$ 1) = 75
• Row 2: (50 $\times$ 0) + (50 $\times$ 0) + (50 $\times$ 0) = 0
• Row 3: (25 $\times$ -1) + (25 $\times$ -1) + (25 $\times$ -1) = -75
• Green Sum: 75 + 0 + -75 = 0

• Row 1: (100 $\times$ 1) + (255 $\times$ 1) + (100 $\times$ 1) = 455
• Row 2: (255 $\times$ 0) + (100 $\times$ 0) + (255 $\times$ 0) = 0
• Row 3: (100 $\times$ 0) + (255 $\times$ 0) + (100 $\times$ 0) = 0
• Blue Sum: 455 + 0 + 0 = 455

Then we add all that we got from the three channels:

The result is the first, top-left pixel (0, 0) of our output feature map.

Then we will slide the kernel and repeat the process until we scanned the whole image. And this is the forward pass.

Backward pass

As we already know, training/learning happens during the backward pass. To help us see how to perform backpropagation in a Convolution layer consider we have these 3 channels and 3 kernels (2 x 2).

If we perform convolution with stride = 1 and padding = 0, we get this output.

Now let’s see how we get the output.

$Y_1 = [ (R_1W_{1,1} + R_2W_{1,2} + R_4W_{1,3} + R_5W_{1,4}) + (B_1W_{2,1} + B_2W_{2,2} + B_4W_{2,3} + B_5W_{2,4}) + (G_1W_{3,1} + G_2W_{3,2} + G_4W_{3,3} + G_5W_{3,4}) ] + b$

$Y_2 = [(R_2W_{1,1} + R_3W_{1,2} + R_5W_{1,3} + R_6W_{1,4}) + (B_2W_{2,1} + B_3W_{2,2} + B_5W_{2,3} + B_6W_{2,4}) + (G_2W_{3,1} + G_3W_{3,2} + G_5W_{3,3} + G_6W_{3,4})] + b$

$Y_3 = [(R_4W_{1,1} + R_5W_{1,2} + R_7W_{1,3} + R_8W_{1,4}) + (B_4W_{2,1} + B_5W_{2,2} + B_7W_{2,3} + B_8W_{2,4}) + (G_4W_{3,1} + G_5W_{3,2} + G_7W_{3,3} + G_8W_{3,4})] + b$

$Y_4 = [(R_5W_{1,1} + R_6W_{1,2} + R_8W_{1,3} + R_9W_{1,4}) + (B_5W_{2,1} + B_6W_{2,2} + B_8W_{2,3} + B_9W_{2,4}) + (G_5W_{3,1} + G_6W_{3,2} + G_8W_{3,3} + G_9W_{3,4})] + b$

In this layer, we need to compute for gradient of the weights $\frac{\partial{L}}{\partial{W}}$ (and biases $\frac{\partial{L}}{\partial{b}})$ which we will use to update/train the model. This value tells us how much the specific weights in our filter contributed to the error.

Then we also need to compute for gradient of the input ($\frac{\partial{L}}{\partial{X}}$) which we will pass backward. In a multi-layer network, the “input” of the current layer was actually the “output” of the previous layer.

• The current layer doesn’t update the previous layer’s weights.
• Instead, it calculates dX to tell the previous layer: “Here is how much you messed up the data you sent me.”
• The previous layer then uses this dX as it’s “incoming error” to calculate its own dW.

Gradient of the learnable parameters

The Chain Rule Foundation

If we have a loss function $L$, an output $y$, and a weight $w$, the relationship looks like this:

1. The weight $w$ affects the output $y$.
2. The output $y$ affects the loss $L$.

To find the gradient of the loss with respect to the weight, we multiply the partial derivatives:

For now, let’s focus on computing the derivative of $W_{1,1}$. If we look at the formulas – ($Y_1$, $Y_2$, $Y_3$, $Y_4$), we see that $W_{1,1}$ is on all four formulas, this means that we need to consider all of them when computing for $\frac{\partial{L}}{\partial{w_{1,1}}}$.

Remember that $\frac{\partial{L}}{\partial{y}}$ is passed backward from the next layer. This is the $dX$ we computed and passed from the next layer.

Then we just have to use the same formula to get the rest of the derivatives of $w$ ($\partial{w_{i,j}}$).

Now the chain rule for $dB$ is:

Since $Y_i = (…) + b$, the derivative of any $Y$ with respect to $b$ is simply 1:

Hence,

Gradient of the input

Now using the same logic we used for $dW$, we calculate $dX$ by looking at every output pixel $y$ that a specific input pixel $x$ contributed to.

So let’s say we want to get $\frac{\partial{L}}{\partial{R_5}}$ (in here we use $R$ just to denote that we are talking about the Red channel, but it really is a pixel from the input $x$). Looking at the formulas above, we can see that $R_5$ was “touched” when we computed for $Y_1$, $Y_2$, $Y_3$, and $Y_4$.

The chain rule for getting $dR_5$ is:

Then we compute for the rest of the $dX$ in a similar fashion.

Implementation

class ConvLayer:

  def __init__(self, in_channels, out_channels, kernel_size=(2, 2), strides=1, padding=1):
    self.strides = strides
    self.padding = padding

    kernel_h, kernel_w = kernel_size

	# He initialization is just another way to set the starting weights before training begins.
	# I will not dive deeper into this in this blog post, but what this does is, it helps our initial values
	# not to become too large or too small which would lead to vanishing/exploding gradients which will prevent
	# us from training our network.
    fan_in = kernel_h * kernel_w * in_channels
    he_init = np.sqrt(2.0 / fan_in)

    self.kernels = np.random.randn(kernel_h, kernel_w, in_channels, out_channels) * he_init
    self.biases = np.zeros(out_channels)

  def forward(self, X):
    self.X = X

    # (batch, height, weight, channel)
    x_n, x_h, x_w, x_ch = X.shape
    kernel_h, kernel_w, _, out_ch = self.kernels.shape

    # Calculate the output size
    y_h = (((x_h - kernel_h) + (2 * self.padding)) // self.strides) + 1
    y_w = (((x_w - kernel_w) + (2 * self.padding)) // self.strides) + 1

    # Create the output
    self.Y = np.zeros((x_n, y_h, y_w, out_ch))

    X = np.pad(X, ((0, 0), (self.padding, self.padding), (self.padding, self.padding), (0, 0)))

    # Slide from top to bottom
    for y_axis in range(0, y_h):
      # Slide from left to right
      for x_axis in range(0, y_w):
        # We need to get the window
        start_x_ax = x_axis * self.strides
        end_x_ax = start_x_ax + kernel_w
        start_y_ax = y_axis * self.strides
        end_y_ax = start_y_ax + kernel_h

        window = X[:, start_y_ax:end_y_ax, start_x_ax:end_x_ax, :]

        for k in range(0, out_ch):
          self.Y[:, y_axis, x_axis, k] = np.sum(window * self.kernels[:, :, :, k], axis=(1, 2, 3)) + self.biases[k]

    return self.Y

  def backward(self, dY):
    # dY shape: (batch, y_h, y_w, out_ch)
    x_n, x_h, x_w, x_ch = self.X.shape
    kernel_h, kernel_w, _, out_ch = self.kernels.shape

    # Initialize gradients
    dX = np.zeros_like(self.X)
    self.dKernels = np.zeros_like(self.kernels)
    self.dBiases = np.zeros_like(self.biases)

    # Pad the input for gradient calculation (consistent with forward pass)
    X_padded = np.pad(self.X, ((0, 0), (self.padding, self.padding), (self.padding, self.padding), (0, 0)), mode='constant', constant_values=0)
    dX_padded = np.zeros_like(X_padded)

    # 1. Bias Gradient: Sum error over batch, height, and width
    # Result shape: (out_channels)
    for k in range(out_ch):
        self.dBiases[k] = np.sum(dY[:, :, :, k])

    # 2. Kernel and Input Gradients
    _, y_h, y_w, _ = dY.shape

    for y_axis in range(y_h):
      for x_axis in range(y_w):
        # Calculate window coordinates
        start_x = x_axis * self.strides
        end_x = start_x + kernel_w
        start_y = y_axis * self.strides
        end_y = start_y + kernel_h

        # Get current window from padded input
        window = X_padded[:, start_y:end_y, start_x:end_x, :]

        for k in range(out_ch):
          # Kernel Gradient: how much the window contributed to the error at this dY position
          # dY[:, y_axis, x_axis, k] is a vector of size (batch,)
          # We need to broadcast it to multiply with the window (batch, kh, kw, in_ch)
          error_pixel = dY[:, y_axis, x_axis, k][:, np.newaxis, np.newaxis, np.newaxis]
          self.dKernels[:, :, :, k] += np.sum(window * error_pixel, axis=0)

          # Input Gradient: distribute the error from this pixel back to the input window
          # Here, we use the weights of the k-th kernel
          dX_padded[:, start_y:end_y, start_x:end_x, :] += error_pixel * self.kernels[:, :, :, k]

    # 3. Remove padding from dX to match original X shape
    if self.padding > 0:
        dX = dX_padded[:, self.padding:-self.padding, self.padding:-self.padding, :]
    else:
        dX = dX_padded

    return dX

  def update(self, learning_rate):
    self.kernels -= learning_rate * self.dKernels
    self.biases -= learning_rate * self.dBiases

Max Pooling

I gave an introduction of what Max Pooling is here. If you don’t know what Max Pooling is, I recommend that you read it to understand what it is, what it does and how it does it (forward pass). I’ll focus on how to do backward pass for a max pool layer here.

Backward pass

During forward pass. max pooling acts like a filter, picking the largest number and ignoring the rest. During backprop, the gradient $dY$ only flows back to the location that had the max value.

For example, consider we have an input $X$:

So passing this to a 2 x 2 max pool with strides = 2. We got this:

To perform backward pass, we need to remember the positions of the max values in the original 4 x 4 matrix. For example, consider we have a gradient coming backwards:

So our $dX$ would be:

Implementation

class MaxPool:
  def forward(self, X, pool_size, stride):
    x_n, x_h, x_w, x_ch = X.shape

    self.X = X
    self.pool_size = pool_size
    self.stride = stride

    # Let's calculate the output sizes
    output_h = ((x_h - pool_size) // stride) + 1
    output_w = ((x_w - pool_size) // stride) + 1
    output = np.zeros((x_n, output_h, output_w, x_ch))

    # Loop through the height
    for h in range(0, output_h):
      # Loop through the width
      for w in range(0, output_w):
        start_x = w * stride
        end_x = start_x + pool_size

        start_y = h * stride
        end_y = start_y + pool_size

        window = np.max(X[:, start_y:end_y, start_x:end_x, :], axis=(1, 2))
        output[:, h, w, :] = window

    return output

  def backward(self, dY):
    # The output will have the same size of the original input
    dX = np.zeros(self.X.shape)

    x_n, _, _, x_ch = self.X.shape # Get batch size and channels from self.X
    _, output_h, output_w, _ = dY.shape # Get output height/width from dY

    # Loop through the height of the output
    for h in range(0, output_h):
      # Loop through the width of the output
      for w in range(0, output_w):
        start_x = w * self.stride
        end_x = start_x + self.pool_size

        start_y = h * self.stride
        end_y = start_y + self.pool_size

        window = self.X[:, start_y:end_y, start_x:end_x, :]
        window_max = np.max(window, axis=(1, 2), keepdims=True)

        mask = (window == window_max) # Shape: (batch_size, pool_size, pool_size, channels)

        # Get the gradient values for the current output cell (h, w)
        # Shape: (batch_size, channels)
        grad_from_next_layer = dY[:, h, w, :]

        # Reshape the gradient to (batch_size, 1, 1, channels) for broadcasting
        # This allows it to be correctly multiplied by the mask (pool_size, pool_size)
        grad_broadcasted = grad_from_next_layer[:, np.newaxis, np.newaxis, :]

        # Add the masked gradient to the corresponding input region in dX
        # Note: using += because pooling regions might conceptually overlap if stride < pool_size,
        # although for max-pooling, only one max element receives the gradient.
        # For typical max pooling (stride == pool_size), += is effectively assigning.
        dX[:, start_y:end_y, start_x:end_x, :] += mask * grad_broadcasted
    return dX

Local Response Normalization (LRN)

Note: While AlexNet used LRN extensively, modern networks largely abandoned it. Practitioners found that a technique called “Batch Normalization” worked better for stability and training speed. But for the sake of learning, we will be using LRN in this blog post.

LRN was formally introduced and popularized by the AlexNet paper (ImageNet Classification with Deep Convolutional Neural Networks).

In 2012, ReLU were new. The authors found that because ReLU neurons have no upper limit (they can output values like 100, 1000, etc.), some neurons in a local neighborhood could become incredibly “loud”, dominating the training process.

Intuition

Imagine a team brainstorming session. If everyone is shouting at the top of their lungs (high activation everywhere), no single good idea stands out. LRN acts as a strict moderator: it forces the loudest person to quiet down their immediate neighbors so their specific idea is heard.

Why do we need LRN?

The necessity of LRN is deeply tied to the activation function: Rectified Linear Unit (ReLU). As you know, this is the formula for the ReLU:

Notice that ReLU is unbounded for positive numbers. If we the output of ReLU to the next layer, it can dominate the gradient calculations during backpropagation. The network becomes overly reliant on a few “shouting” neurons.

By applying LRN, we bound these raw activations relatively. We ensure that a neuron is only considered highly active if it is significantly more active than the filters adjacent to it in the channel stack.

Here’s the formula to compute the normalized (LRN) output:

Where:

• $b_{x,y}^i$: The normalized output.
• $a_{x,y}^i$: The original, un-normalized activation of a neuron at spatial location $(x, y)$ in the $i$-th channel.
• $N$: The total number of channels (kernels) in that layer.
• $n$: The size of the local neighborhood (how many adjacent channels we look at).
• $k, \alpha, \beta$: Hyperparameters. $k$ prevents division by zero, $\alpha$ is a scaling factor, and $\beta$ controls the strength of the contrast (the normalization effect).

Example

Imagine we have this as the output of ReLU.

Now let’s start looking at $a^{1}_{0,0} = 32$. Then we decide n which is how many adjacent channels we look at. The most intimidating part of the formula is $\sum_{j=\max(0, i – n/2)}^{\min(N-1, i + n/2)} (a_{x,y}^j)^2)$. But it can be broken down into a few steps:

1. First we want to get the “radius”. The formula is $\text{floor}(n/2)$. So if $n = 3$, then the radius is $\text{floor}(3/2) = 1$.
2. Then the channels we will include is $i$ which is the current channel. Then since our radius is 1 then we will include 1 channel at the top and 1 channel at the bottom of the current channel $i$.
3. So for example, if we want to get the LRN of $a^{2}_{0,0} = 29$ and we choose $n = 3$, then we look at the layer above it $a^{1}_{0,0} = 32$ and at the layer below it $a^{3}_{0,0} = 39$.
4. Then we will get the squared sum of these values. So $(a^{2}_{0,0} = 29)^2 + (a^{1}_{0,0} = 32)^2 + (a^{3}_{0,0} = 39)^2 = 841 + 1024 + 1521 = 3386$

To complete the full formula:

Let use the following hyperparameters.

• $k$: 1
• $\alpha$: 0.0001
• $\beta$: 0.75

Note: In our example, we used the value in the middle layer so we got a layer above and its bottom. But what if we want to get $a^{1}_{0,0}$? Since this is the top layer (there’s no layer above it), then we will only include $a^{1}_{0,0}$ and the layer below it, $a^{2}_{0,0}$.

Backward pass

In the LRN layer, we don’t have any learnable parameters (weights and biases), we only need to compute for $dX$.

Remember that because of the sliding window in LRN, a single input $x_i$ influences multiple outputs in its neighborhood $y_j$.

Let’s look at the LRN formula again.

So to find the total gradient for $x_i$, we must add up the gradients flowing backwards from every output it helped calculate.

Basically we want to get the gradient of $y_i$ with respect to $x_i$ when $i = j$ and add it to gradient of $y_j$ with respect to $x_i$ when $i \neq j$.

To make derivation easier for us, let us set the denominator:

So we have the simplified formula:

Now remember the quotient rule:

Let: $u = x_i$ and $v = (s_i)^\beta$

$v’ = \frac{\partial}{\partial{x_i}}(s_i^\beta)$, so using the chain rule

Outside: $\beta (s_i)^{\beta-1}$
Inside: $s_i = k + \alpha\sum_j x_j^2$ -> $\frac{\partial}{\partial{x_i}} = 2\alpha x_i$ All other $x$ values where $i \neq j$ are treated as constants.

So for $i = j$:

Then we get the gradient where $i \neq j$

Outside: $\beta (s_j^{-\beta – 1})$
Inside: $2\alpha x_i$

Combining them all together

The first term is the upstream gradient at exact pixel & channel $i$.
The second term is the gradient routed backward directly from its own channel ($i = j$)
The third term is the sum of all the extra gradients routed from the neighboring channels that $x_i$ helped normalize ($i \neq j$)

Lateral Inhibition

What we just did is called Lateral Inhibition. Lateral Inhibition is a concept borrowed directly from neurobiology and applied to computer vision to make features stand out more sharply.

The biological inspiration is our eye. In our retina, we have millions of photoreceptor cells. When we look at a sharp edge (like a black line on white paper), the cells receiving the bright light fire intensely.

However, if every cell just fired based on what it saw, the image might look blurry or “fuzzy” around the edges. To fix this, nature evolved Lateral Inhibition.

What happens is when a neuron fires excitedly, it sends a chemical signal to its immediate neighbors telling them to “shut up” (reduce their activity). The result is, the strong signal stays strong but the surrounding “noise” gets muted.

We mimic Lateral Inhibition with LRN. In our network, the “neurons” are the different filters looking at the same pixel location.

The Situation: Imagine you have 5 different filters looking at a picture of a cat’s ear.

• Filter A (Vertical Edge) says: “I see a little bit of an edge here.” (Value: 50)
• Filter B (Curve Detector) says: “I see a PERFECT curve here!” (Value: 200)
• Filter C (Color Blob) says: “I see some brown.” (Value: 20)

Without LRN: All three values (50, 200, 20) pass to the next layer. The signal is “muddy.”

With LRN (Lateral Inhibition): The math looks at the neighbors. It sees that Filter B is screaming with a value of 200.

• It uses Filter B’s high value to divide (suppress) the values of Filter A and C.
• Filter A gets crushed down to near zero.
• Filter B survives (mostly) intact.

Implementation

class LRN:
  def forward(self, X, neighbors=5, k=2, alpha=1e-4, beta=0.75):
    # X shape is (batch, height, width, channels)
    self.X = X
    self.alpha = alpha
    self.beta = beta
    self.neighbors = neighbors
    self.k = k

    radius = neighbors // 2

    # (batch, height, width, channel)
    x_n, x_h, x_w, x_ch = X.shape

    # Output is the same size as input.
    output = np.zeros(X.shape)

    # Cache for the backward pass
    self.scale = np.zeros(X.shape)

    # Loop through the channels
    for c in range(0, x_ch):
      a_i = X[:, :, :, c]

      top = max(0, c - radius)
      bottom = min(x_ch, c + radius + 1)
      chunk = X[:, :, :, top:bottom]

      S_c = k + (alpha * np.sum(np.square(chunk), axis=3))
      self.scale[:, :, :, c] = S_c

      output[:, :, :, c] = a_i / (S_c ** beta)

    self.Y = output

    return output

  def backward(self, dY):
    dX = np.zeros_like(self.X)
    radius = self.neighbors // 2
    x_ch = self.X.shape[3]

    # Precompute the common term used in the neighborhood summation
    # dY_j * (Y_j / S_j)
    sum_term = dY * (self.Y / self.scale)

    for c in range(x_ch):
      top = max(0, c - radius)
      bottom = min(x_ch, c + radius + 1)

      # Sum the precomputed term over the neighborhood window
      chunk_sum = np.sum(sum_term[:, :, :, top:bottom], axis=3)

      # Part 1: Direct effect
      term1 = dY[:, :, :, c] / (self.scale[:, :, :, c] ** self.beta)

      # Part 2: Indirect effect from neighbors
      term2 = 2 * self.alpha * self.beta * self.X[:, :, :, c] * chunk_sum

      # Combine
      dX[:, :, :, c] = term1 - term2

    return dX

Dense (Fully Connected) Layers

For the dense (fully connected) layers, it is simply the linear operation Y = x \cdot w + b then followed by ReLU for hidden layers, then on the final classification layer we use the Softmax function. I recommend that you read “Building an Image Classifier from Scratch“.

Dropout

Intuition

Imagine you manage a team of engineers working on a critical project. It is typical that one engineer, let’s name her “Alice” becomes a “superstar” of the project. She is so good that everone else in the team starts relying on her. If Alice solves a problem, “Bob” and “Charlie” don’t bother learning how to solve it; they just wait for Alice’s output. In machine learning terms, Bob and Charlie have co-adapted to Alice.

The problem arises when Alice gets sick (or when our model encounters new data that looks slightly different). The team falls apart because no one else knows what to do. The team has overfitted to Alice’s presence.

As a good manager, you thought of a solution to this issue. You randomly tell your engineers to stay home each day.
– On Monday, Alice should not come to the office and work. Now, Bob and Charlie must learn to do Alice’s job to get the project done.
– On Tuesday, it’s Bob’s turn to take the day’s off. Alice and Charlie must cover for him.

By forcing different subsets of team to work together, you ensure that no single member (neuron) becomes too important. Every neuron must be capable of extracting useful features independently of who else is present. This creates a resilient “team” (network) that generalized better to unseen problems.

How does this help build better models?

1. Regularization: By randomly disabling neurons, we prevent the network from memorizing the noise in the training data. The network cannot rely on just the complex combinations of neurons to solve the problem because those combinations will be different on each dropout training. This reduces Overfitting.
2. Approximate Model Ensembling
   • A network with n neurons has 2n possible sub-networks (different combinations of dropped/kept neurons).
   • Every time we pass a batch of data through the network during training, we are training a different random sub-network.
   • However, all these sub-networks share the same weights W.

Training with Dropout can be seen as training a massive ensemble of exponentially many different neural networks that share parameters. At test time, when all of the neurons are active (no dropout), we are effectively averaging the predictions of this massive ensemble. Ensemble almost always outperform single models.

Implementation

class Dropout:
  def __init__(self, rate=0.5):
    # rate is the percentage of neurons to turn off (0.5 = 50%)
    self.rate = rate
    self.keep_prob = 1.0 - rate
    self.mask = None

  def forward(self, X, training=True):
    if training:
      # Create a random mask of 0s and 1s, scaled by keep_prob
      self.mask = (np.random.rand(*X.shape) < self.keep_prob) / self.keep_prob
      return X * self.mask
    else:
      # If not training, let all data pass through normally
      return X

  def backward(self, dY):
    # Only pass gradients back through the neurons we kept active
    return dY * self.mask

Mini AlexNet

At this point, we should have a working understanding of the layers that builds the AlexNet architecture. So let’s build a smaller, more simplified version of it and call it Mini Alexnet. The purpose of this is so we can see how we the different layers we learned above work together to create a deep neural network for image classification.

A few things to keep in mind:

1. We will be using Cifar10 dataset. The dataset contains 60,000 (32 x 32) colored images in 10 classes (classifications).
2. Our implementation only uses a portion of this data so that we can see it train. Our code is not optimized for training. Our goal is to see and understand the inner workings of the AlexNet architecture and using the optimized approach will make it harder for us to achieve this goal.

Loading Cifar10 dataset

import tensorflow as tf
from tensorflow.keras.datasets import cifar10

(x_train, y_train), (x_test, y_test) = cifar10.load_data()

ReLU, Dense, SoftMax, CCE

I didn’t include these layers and their implementations in this blog post since we already covered it in my previous posts (ReLU, Dense, SoftMax and CCE), but for the sake of completeness, here’s the code implementation.

class ReLU:
  def forward(self, X):
    self.X = X
    return np.maximum(0, X)

  def backward(self, dY):
    self.dX = dY.copy()
    self.dX[self.X <= 0] = 0
    return self.dX

class Dense:
  def __init__(self, input_count, output_count):
    # Calculate fan_in for He Initialization
    fan_in = input_count
    he_init = np.sqrt(2.0 / fan_in)

    self.W = np.random.randn(input_count, output_count) * he_init
    self.b = np.zeros(output_count)

  def forward(self, X):
    # X has the shape (batch, number of features)
    self.X = X
    return np.dot(X, self.W) + self.b

  def backward(self, dY):
    # We need to compute for three things.
    self.dW = np.dot(self.X.T, dY)
    self.dB = np.sum(dY, axis=0)

    return np.dot(dY, self.W.T)

  def update(self, learning_rate):
    self.W -= learning_rate * self.dW
    self.b -= learning_rate * self.dB

class SoftMax():
  def forward(self, X):
    # The shape of X is (batch, input_features)
    self.X = X

    X_shift = X - np.max(X, axis=1, keepdims=True)
    exps = np.exp(X_shift)

    self.output = exps / np.sum(exps, axis=1, keepdims=True)

    return self.output

class CategoricalCrossEntropy:
  def forward(self, y_pred, y_true, epsilon=1e-8):
    # Epsilon to prevent infinity.
    mult = y_true * np.log(y_pred + epsilon)
    sum = - np.sum(mult, axis=1, keepdims=True)
    return np.mean(sum)

  def backward(self, y_pred, y_true):
    self.back_output = (1/len(y_true)) * (y_pred - y_true)
    return self.back_output

Mini AlexNet

import numpy as np

# Add a seed for reproducibility
np.random.seed(42)

class MiniAlexNet:
  def __init__(self):
    self.data = []

    # Initialize the layers.
    self.conv_1 = ConvLayer(in_channels=3, out_channels=96, kernel_size=(3, 3), strides=1, padding=1)
    self.relu_1 = ReLU()
    self.lrn_1 = LRN()
    self.max_1 = MaxPool()

    # Conv 2.
    self.conv_2 = ConvLayer(in_channels=96, out_channels=256, kernel_size=(3, 3), strides=1, padding=1)
    self.relu_2 = ReLU()
    self.lrn_2 = LRN()
    self.max_2 = MaxPool()

    ## Conv 3.
    self.conv_3 = ConvLayer(in_channels=256, out_channels=384, kernel_size=(3, 3), strides=1, padding=1)
    self.relu_3 = ReLU()

    ## Conv 4.
    self.conv_4 = ConvLayer(in_channels=384, out_channels=384, kernel_size=(3, 3), strides=1, padding=1)
    self.relu_4 = ReLU()

    ## Conv 5
    self.conv_5 = ConvLayer(in_channels=384, out_channels=256, kernel_size=(3, 3), strides=1, padding=1)
    self.relu_5 = ReLU()
    self.max_3 = MaxPool()

    # Calculate the flattened size for a 32x32 input image: 4 * 4 * 256 = 4096
    flattened_size = 4096

    # Initialize the Dense layers ONCE
    self.dense_1 = Dense(input_count=flattened_size, output_count=1024)
    self.relu_dense_1 = ReLU()
    self.dropout_1 = Dropout(rate=0.5)

    self.dense_2 = Dense(input_count=1024, output_count=1024)
    self.relu_dense_2 = ReLU()
    self.dropout_2 = Dropout(rate=0.5)

    self.dense_3 = Dense(input_count=1024, output_count=10) # 10 classes
    self.softmax = SoftMax()

  def forward(self, X, y, training=False):
    # Conv 1
    out = self.conv_1.forward(X)
    out = self.relu_1.forward(out)
    out = self.lrn_1.forward(out)
    out = self.max_1.forward(out, pool_size=2, stride=2)

    ## Conv 2
    out = self.conv_2.forward(out)
    out = self.relu_2.forward(out)
    out = self.lrn_2.forward(out)
    out = self.max_2.forward(out, pool_size=2, stride=2)

    ## Conv 3
    out = self.conv_3.forward(out)
    out = self.relu_3.forward(out)

    ## Conv 4
    out = self.conv_4.forward(out)
    out = self.relu_4.forward(out)

    ## Conv 5
    out = self.conv_5.forward(out)
    out = self.relu_5.forward(out)
    out = self.max_3.forward(out, pool_size=2, stride=2)

    # Remember the shape before flatten
    self.pre_flatten_shape = out.shape

    # Flatten
    flattened = out.reshape(out.shape[0], -1)

    # Dense 1
    out = self.dense_1.forward(flattened)
    out = self.relu_dense_1.forward(out)
    out = self.dropout_1.forward(out, training=training)

    # Dense 2
    out = self.dense_2.forward(out)
    out = self.relu_dense_2.forward(out)
    out = self.dropout_2.forward(out, training=training)

    # Dense 3
    out = self.dense_3.forward(out)

    # Softmax
    softmax_output = self.softmax.forward(out)

    return softmax_output

  def train(self, X, y, epochs, batch_size=32, learning_rate=0.001):
    self.cce = CategoricalCrossEntropy()
    num_samples = X.shape[0]

    for epoch in range(epochs):
      print(f"--- Starting Epoch {epoch + 1}/{epochs} ---")

      # We shuffle an array of indices, then apply it to both X and y so they stay matched
      indices = np.arange(num_samples)
      np.random.shuffle(indices)
      X_shuffled = X[indices]
      y_shuffled = y[indices]

      epoch_loss = 0

      # 2. Iterate through the dataset in mini-batches
      for start_idx in range(0, num_samples, batch_size):
        # Calculate the end index (and handle the leftover remainder at the end)
        end_idx = min(start_idx + batch_size, num_samples)

        # 3. Slice the batch
        X_batch = X_shuffled[start_idx:end_idx]
        y_batch = y_shuffled[start_idx:end_idx]

        # Forward pass
        output = self.forward(X_batch, y_batch, training=True)

        batch_loss = self.cce.forward(output, y_batch)

        # Backpass
        out = self.cce.backward(output, y_batch)

        # Dense 3
        out = self.dense_3.backward(out)

        # Dense 2
        out = self.dropout_2.backward(out)
        out = self.relu_dense_2.backward(out)
        out = self.dense_2.backward(out)

        # Dense 1
        out = self.dropout_1.backward(out)
        out = self.relu_dense_1.backward(out)
        out = self.dense_1.backward(out)

        # Unflatten
        out = out.reshape(self.pre_flatten_shape)

        # Conv 5
        out = self.max_3.backward(out)
        out = self.relu_5.backward(out)
        out = self.conv_5.backward(out)

        # Conv 4
        out = self.relu_4.backward(out)
        out = self.conv_4.backward(out)

        # Conv 3
        out = self.relu_3.backward(out)
        out = self.conv_3.backward(out)

        # Conv 2
        out = self.max_2.backward(out)
        out = self.lrn_2.backward(out)
        out = self.relu_2.backward(out)
        out = self.conv_2.backward(out)

        # Conv 1
        out = self.max_1.backward(out)
        out = self.lrn_1.backward(out)
        out = self.relu_1.backward(out)
        out = self.conv_1.backward(out)

        ## Then we update the parameters
        self.dense_3.update(learning_rate)
        self.dense_2.update(learning_rate)
        self.dense_1.update(learning_rate)

        self.conv_5.update(learning_rate)
        self.conv_4.update(learning_rate)
        self.conv_3.update(learning_rate)
        self.conv_2.update(learning_rate)
        self.conv_1.update(learning_rate)

        # Batch loss
        print(f"Batch: {start_idx} - {end_idx}, Loss: {batch_loss}")
        epoch_loss += batch_loss

      # Print the average loss for the entire epoch
      num_batches = np.ceil(num_samples / batch_size)
      avg_loss = epoch_loss / num_batches
      print(f"Epoch {epoch + 1}, Avg Loss: {avg_loss:.4f}")
      # Append the avg loss
      self.data.append(avg_loss)

# The shape of x_train is (Number of images, height, width, channel)
# Divide by 255.0 to scale the pixels between 0 and 1
x_sample = x_train[:100].astype(np.float32) / 255.0

# Convert the labels into One-Hot Encoding format for the math to work
num_classes = 10
y_sample = y_train[:100]

y_sample_flatten = y_sample.flatten()

y_sample_one_hot = np.zeros((y_sample_flatten.size, num_classes))
y_sample_one_hot[np.arange(y_sample_flatten.size), y_sample_flatten] = 1

mini_alex_net = MiniAlexNet()
mini_alex_net.train(x_sample, y_sample_one_hot, epochs=200)

As mentioned above, the code above is not really optimized for production. But you can still run it, here’s how it looks like. Take note that I severely limited the resources (number of samples = 100, batch size = 32) because I just want to see how it looks like when it rains. And as you can see on the avg loss per epoch, while it goes up and down, you can see the trend going down which is typically what we want when we are training model.

What’s Next?

I hope you enjoyed reading and learning as much as I did! And as mentioned throughout this post, our implementation of the layers are not optimized for training (but for learning!). So it’s a good learning opportunity to do them!

• Vectorize the operations. This will save a whole lot of computations because you don’t have to do the computations in a loop.
• Look into im2col technique.
• Research about different kinds of initialization. In this blog post we used He Initialization but there are some others: Xavier/Glorot init, LeCun init, etc.
• After optimizing your implementation, try implementing the full version of AlexNet.