Flatten Layer - CNNs & Computer Vision

Flatten Layer

What you'll build

The layer that turns a batch of feature maps into a batch of vectors so it can feed a fully connected (linear) layer.

Background

A CNN ends with feature maps of shape (batch, height, width, channels), but a linear classifier expects (batch, features). Flattening keeps the batch axis and concatenates everything else. It has no parameters; it only changes the view of the data.

Layout used here: NHWC

The input is (N, H, W, C) = (batch, height, width, channels), the channels-last layout of TensorFlow/Keras. PyTorch's nn.Flatten usually sees NCHW (N, C, H, W) instead. Both flatten in row-major (C) order, so the element order of the output differs between layouts: weights trained on one layout don't transfer to the other without a permute.

Formula

out[n, (h * W + w) * C + c] = X[n, h, w, c]        out shape: (N, H * W * C)

That is exactly NumPy's default row-major reshape.

Inputs

• X: NumPy array of shape (N, H, W, C), any dtype. It may be non-contiguous (for example a transposed view).

Expected return format

A NumPy array of shape (N, H * W * C) with the same values in row-major order and the same dtype.

Edge cases and conventions

• N = 1 must still return a 2D array (1, H*W*C), not a 1D vector. Never flatten the batch axis.
• Use -1 to let NumPy infer the feature size.

Constraints

NumPy only, no loops needed.

Follow-up

• Why do modern CNNs replace flatten + a huge linear layer with global average pooling?
• In PyTorch, why can x.view(N, -1) fail on a permuted tensor while x.reshape(N, -1) works?

Input: X.shape = (2, 4, 4, 3)
Output: shape (2, 48)

Input: X = np.arange(12).reshape(2, 3, 2, 1)
Output: [[0, 1, 2, 3, 4, 5], [6, 7, 8, 9, 10, 11]]

import numpy as np

def flatten(X):
    """
    Flatten NHWC feature maps to vectors.

    Args:
        X: Input array, shape (batch, height, width, channels)

    Returns:
        output: Array of shape (batch, height * width * channels), row-major order
    """
    # Your code here
    pass