Lesson 3: Working with large data — NumPy & Pandas in depth

In Lesson 2 we understood the mathematics of linear algebra and how to optimise along a slope. But when training real AI models, the input is not a handful of numbers — it is millions of pixels, billions of words, or enormous database tables.

Keep using Python's default for loops and lists on data at that scale and your program will run unacceptably slowly. This lesson explains the root cause at the hardware level, and how to fix it properly by making use of AI's two foundational libraries: NumPy and Pandas.

3.1 Why are Python for loops so slow?

As Lesson 1 said, Python is an interpreted, dynamically typed language — it runs directly with no compile step, and a variable can hold a number on one line and a string on the next. Both traits are convenient to write, but they cost something, and this section shows where that cost sits.

When you create an ordinary list my_list = [1, 2, 3], Python does not store those three numbers next to each other in RAM. It creates an array of pointers — each pointer being an address, pointing at a full Python object (called a PyObject) scattered somewhere in memory.

Every time you run a loop for x in my_list, on each iteration the Python interpreter is forced to do a series of things you never wrote — the term for which is overhead:

• Type checking: check whether x on this iteration is an integer, a float or a string, to decide which operation applies.
• Boxing/unboxing: extract the raw numeric value from deep inside the complex PyObject structure.
• Memory management: update the reference count so garbage collection works.

3.2 The principle of vectorization

Vectorization is the technique of replacing explicit loops with operations applied directly to a whole array. Instead of writing a loop that handles one element at a time in the slow Python layer, we push the entire looping job down into NumPy's high-performance machine-code layer.

On top of that, every modern CPU supports a hardware-level parallel instruction set called SIMD (Single Instruction, Multiple Data). This lets the CPU apply one operation — a multiplication, say — to several memory slots at once, within a single clock cycle.

Let's compare the performance of a pure Python loop against a vectorised NumPy operation:

import time
import numpy as np

size = 1_000_000
python_list = list(range(size))
numpy_array = np.arange(size)

# Way 1: a plain Python loop (a list comprehension is still a Python loop).
start = time.time()
python_result = [x * 2 for x in python_list]
print(f"Python loop:        {(time.time() - start) * 1000:7.2f} ms")

# Way 2: one vectorised NumPy operation. The loop still happens — but inside
# compiled C, over a contiguous block, using the CPU's SIMD instructions.
start = time.time()
numpy_result = numpy_array * 2
print(f"NumPy vectorised:   {(time.time() - start) * 1000:7.2f} ms")

Results on the machine this was written on (Apple M1 Max, 32 GB RAM, Python 3.11, NumPy 2.4):

Python loop:          22.14 ms
NumPy vectorised:      0.39 ms

Roughly 20 to 56 times faster — I ran it repeatedly and the ratio moved around inside that range, with the first run always the worst because the CPU cache is still cold. Don't attach yourself to any one figure: the ratio depends on the machine, on the data type, and on the operation itself. What is worth remembering is the order of magnitude — tens of times, not a few percent.

One thing that is easy to misread also needs saying plainly: the loop does not disappear. A million multiplications still have to happen. They simply move from the Python layer down into compiled C, where there is no type check and no unboxing on each step. "Vectorising" means moving the loop somewhere else, not deleting it.

3.3 Broadcasting in NumPy

In standard linear algebra you can only add two matrices of exactly the same size. In machine learning, though, we constantly need operations across mismatched shapes — adding a bias vector to every row of a data matrix, for example.

NumPy supports an extremely powerful mechanism for this called broadcasting. It automatically aligns and virtually stretches the smaller array's dimensions to match the larger one during the computation, without copying any real data in RAM.

Two arrays can be broadcast together if the sizes of their dimensions satisfy one of these rules (working from right to left):

1. The dimension is the same size in both arrays.
2. One of the arrays has size 1 in that dimension (which then gets virtually stretched to match).
3. The dimension does not exist in the array with fewer dimensions. When two arrays differ in number of dimensions, NumPy prepends dimensions of size 1 to the front of the shorter one until they match, and only then applies the two rules above. This is why a vector (3,) is read as (1, 3) and not (3, 1) — and it is exactly where the pitfall at the end of this section comes from.

The phrase "from right to left" is the most important part of those three rules. NumPy aligns two shapes on their right edge, the way you align two decimal numbers on the decimal point rather than on their first digit.

An illustration of broadcasting:

When we compute $A + B$, NumPy sees that the last dimensions are equal ($3 = 3$) and that $B$'s first dimension is $1$. It automatically "virtually duplicates" the vector $B$ down the row dimension to form a $2 \times 3$ matrix, and adds directly:

import numpy as np

A = np.array([[1, 2, 3],
              [4, 5, 6]])          # shape (2, 3)

B = np.array([10, 20, 30])         # shape (3,) — read as (1, 3) when broadcasting

# Adding straight away: B is repeated down the rows, virtually.
result = A + B
print("result:\n", result)
# result:
#  [[11 22 33]
#   [14 25 36]]

import numpy as np

matrix = np.arange(12).reshape(4, 3)   # shape (4, 3) — 4 rows, 3 columns
row_bias = np.array([1, 2, 3, 4])      # shape (4,)  — one value per ROW

try:
    result = matrix + row_bias
except ValueError as e:
    print("error:", e)
    # error: operands could not be broadcast together with shapes (4,3) (4,)
    # NumPy aligns from the RIGHT: 3 (matrix) against 4 (row_bias) -> no match.

# The fix: add a virtual dimension, turning (4,) into (4, 1).
row_bias_col = row_bias[:, np.newaxis]   # same as row_bias.reshape(4, 1)
result = matrix + row_bias_col           # (4,3) + (4,1) -> valid, one value per row
print("added per row:\n", result)
# added per row:
#  [[ 1  2  3]
#   [ 5  6  7]
#   [ 9 10 11]
#   [13 14 15]]

The rule to remember: a flat (N,) vector broadcasts across COLUMNS by default (each element belonging to one column, repeated down every row). To broadcast down ROWS you must deliberately turn it into a column vector (N, 1) first, using reshape or np.newaxis.

3.4 Pandas: tables with named columns

The three sections above were all about NumPy, and NumPy has a precondition we haven't mentioned: every cell in an array must be the same type. That's fine for images and weight matrices. But real data is rarely that tidy — a CSV exported from some system typically has a name column (text), a date column, a numeric column, and empty cells scattered through it. Push that table straight into NumPy and everything gets coerced to strings, and you lose the ability to compute at all.

That is where Pandas comes in.

DataFrame and Series — the only two types you need to know

A DataFrame is a two-dimensional table with named columns, each column having its own data type. An everyday way to picture it: it is an Excel sheet you drive with code. It has column headers, a number of rows, each column holds one kind of information, and empty cells are normal.

A Series is a single column taken out of that table. This distinction matters more than it looks, because how you write your brackets decides which one you get back — and that is the first pitfall below.

Two examples to show the boundary between Pandas and NumPy.

• Pandas' job: a 50,000-row table of orders with an order ID, customer name, order date and amount, where 3% of rows are missing the date. You need to filter by month, group by customer, fill the gaps. Named columns, mixed types, empty cells — Pandas.
• NumPy's job: a batch of 50,000 images, each 224×224×3 pixels, all numbers. There is no meaningful "column name" here, and you need operations that run fast over 7.5 billion numbers — NumPy. Stuffing this block into a DataFrame only makes it slower and harder to use.

Build a table and look at it before doing anything else

We'll use one small score table as the running example for this section. It deliberately contains two empty cells, because real data always has empty cells and handling them is most of the actual work:

import numpy as np
import pandas as pd

df = pd.DataFrame(
    {
        "student_id": [1, 2, 3, 4, 5, 6],
        "name": ["An", "Binh", "Chi", "Dung", "Em", "Phuc"],
        "class": ["A", "A", "B", "B", "A", "B"],
        "math": [8.5, 6.0, 9.0, np.nan, 7.5, 5.5],       # np.nan = an empty cell
        "literature": [7.0, 8.5, 6.5, 7.0, np.nan, 6.0],
    }
)

print("shape:", df.shape)      # (rows, columns)
print(df.head(3))              # the first 3 rows — always look before you compute
print(df.isna().sum())         # how many empty cells in each column

shape: (6, 5)
   student_id  name class  math  literature
0           1    An     A   8.5         7.0
1           2  Binh     A   6.0         8.5
2           3   Chi     B   9.0         6.5

student_id    0
name          0
class         0
math          1
literature    1
dtype: int64

df.isna().sum() deserves to become your first reflex on opening any new dataset. It tells you how many cells each column is missing — here 1 in math and 1 in literature. Skip this step and those empty cells will quietly turn every average you compute later into NaN.

Selecting data: [] vs .loc vs .iloc

This is the single most confusing thing for people new to Pandas, so it gets its own section. There are three ways to select and they are not interchangeable:

• df["math"] — selects one column by name and returns a Series. Whereas df[["math"]], with two sets of brackets, returns a DataFrame containing one column. Same data, different type — and this is the cause of a great many baffling AttributeErrors later on.
• df.loc[...] — selects by label: column names, and filter conditions. This is what you'll use 90% of the time.
• df.iloc[...] — selects by numeric position, exactly like slicing a Python list. The "i" is for integer. Use it when you want "the first 3 rows" regardless of what they're called.

print(type(df["math"]).__name__)      # Series    — one set of brackets
print(type(df[["math"]]).__name__)    # DataFrame — two sets of brackets

# .loc takes a condition and a list of columns. Read it as a sentence:
# "rows where math is above 7, and only the name and math columns".
print(df.loc[df["math"] > 7, ["name", "math"]])

# .iloc counts positions instead, exactly like list slicing.
print(df.iloc[0:2, 1:3])              # first 2 rows, columns at positions 1 and 2

Series
DataFrame
  name  math
0   An   8.5
2  Chi   9.0
4   Em   7.5
   name class
0    An     A
1  Binh     A

Notice the result of .loc: the surviving row numbers are 0, 2, 4 — with gaps. Those are row labels, carried over from the original table rather than renumbered from 0. This detail causes surprise constantly: after filtering, df.iloc[1] and df.loc[1] return different rows. Internalising the label-versus-position distinction here means you have got past the hardest part of Pandas.

Handling empty cells

There are two options, and choosing the wrong one damages your data:

• dropna() — discard rows containing an empty cell. Tidy, but if 3% of rows have a gap you lose 3% of your data, and what you lose is usually not lost at random.
• fillna(value) — fill the gap. For numeric columns, filling with that column's own mean or median is the most common approach.

clean = df.copy()   # never modify the original while exploring

# Fill each numeric column with ITS OWN mean, not one mean for the whole table.
clean["math"] = clean["math"].fillna(clean["math"].mean())
clean["literature"] = clean["literature"].fillna(clean["literature"].mean())

print(clean[["name", "math", "literature"]])

   name  math  literature
0    An   8.5         7.0
1  Binh   6.0         8.5
2   Chi   9.0         6.5
3  Dung   7.3         7.0
4    Em   7.5         7.0
5  Phuc   5.5         6.0

Dung's row receives 7.3 in the math column: the mean of the 5 values that do exist, $(8.5 + 6.0 + 9.0 + 7.5 + 5.5) / 5 = 7.3$. Worth noticing that mean() skips empty cells automatically — it divided by 5, not 6. That default is convenient but it can also stop you noticing your table has gaps, which is one more reason to run isna().sum() first.

groupby — the thing that makes learning Pandas worth it

If there were only one reason to use Pandas instead of writing your own code, it would be groupby. It gathers rows by the value of a column and then computes on each group — a job that takes twenty error-prone lines with a loop and a dict.

# Read it as: split the rows by class, then average these two columns in each group.
print(clean.groupby("class")[["math", "literature"]].mean().round(3))

        math  literature
class
A      7.333         7.5
B      7.267         6.5

One line, and you have each class's average in each subject. Notice the class column has become the row label of the result — that is groupby's default behaviour; add as_index=False if you would rather keep it an ordinary column.

Handing over to NumPy — where the two halves of this lesson meet

AI models do not accept DataFrames. The final step of every preprocessing pipeline is converting the cleaned table into a pure numeric array — and that is precisely where this lesson's two halves join:

# Pick only the numeric columns — a model has no use for the name column.
features = clean[["math", "literature"]].to_numpy()

print(type(features).__name__, features.shape, features.dtype)
# ndarray (6, 2) float64

# From here on it is pure NumPy again: broadcasting works, vectorisation works.
print(features - features.mean(axis=0))   # centre each column on its own mean

.to_numpy() is the answer to "so where do these two libraries connect". Pandas takes the real world's messy data and hands out a clean block of numbers; NumPy takes that block and computes. Note axis=0 on the last line: it means "collapse along the row dimension", giving the mean of each column. Remembering that axis=0 is down columns and axis=1 is across rows will save you a great deal of trial and error.

3.5 Hands-on project: preprocessing a batch of images with both libraries

This lesson's project joins the exact chain described at the top: Pandas inspects and cleans, NumPy computes. The program generates a batch of fake grayscale images (a 3-dimensional Batch × Height × Width array), builds a per-image statistics table with Pandas, uses that table to detect and drop the broken images, then normalises what remains with NumPy — with no loop running over pixels at all.

Two points about how the fake data is built are worth stating first, because they decide whether the project teaches anything. One: each image gets its own brightness level (40, 55, 70…) rather than all coming from the same random distribution — if every image were statistically identical, the Pandas table in the middle would be full of identical numbers and tell you nothing. Two: two images are broken on purpose, one pitch black and one pure white, so the cleaning step has real work to do. Real data always contains frames like these — a dead sensor, or an overexposure.

The full source:

# image_normalize.py
# Lesson 3: Working with large data — NumPy & Pandas in depth
# Practical AI Engineer series
#
# Run it with:  python image_normalize.py
# Requires:     pip install numpy pandas
#
# A small but complete preprocessing pipeline, the shape you meet in real
# projects: Pandas inspects and cleans the batch metadata, then NumPy does the
# heavy per-pixel arithmetic. Not one `for` loop over pixels anywhere.

import numpy as np
import pandas as pd

BATCH, HEIGHT, WIDTH = 12, 28, 28


def generate_dummy_images(num_images=BATCH, height=HEIGHT, width=WIDTH):
    """Fake a batch of grayscale images that genuinely differ from each other.

    Each image gets its own brightness level and its own amount of noise, so the
    per-image statistics below actually vary. (Drawing every image from the same
    uniform 0-255 distribution would make them statistically identical, and the
    normalisation step would have nothing to show.)

    Two images are deliberately broken, to give the cleaning step real work:
    one is entirely black, one is entirely white.
    """
    rng = np.random.default_rng(42)
    images = np.empty((num_images, height, width), dtype=np.uint8)

    for i in range(num_images):  # per IMAGE, not per pixel — 12 iterations, not 9408
        brightness = 40 + i * 15  # 40, 55, 70, ... a different level each time
        noise = rng.normal(0.0, 18.0, size=(height, width))
        images[i] = np.clip(brightness + noise, 0, 255).astype(np.uint8)

    images[3] = 0  # a dead sensor: completely black
    images[8] = 255  # an overexposed frame: completely white
    return images


def describe_batch(images):
    """Build a Pandas table of per-image statistics — one row per image.

    This is what Pandas is for: a small table with named, mixed-type columns that
    you want to inspect, filter and group. The pixels themselves stay in NumPy.
    """
    flat = images.reshape(len(images), -1)  # (batch, height*width), still no loop
    df = pd.DataFrame(
        {
            "image_id": np.arange(len(images)),
            "camera": ["cam-a", "cam-b"] * (len(images) // 2),
            "mean": flat.mean(axis=1),  # axis=1 -> collapse pixels, keep images
            "std": flat.std(axis=1),
            "min": flat.min(axis=1),
            "max": flat.max(axis=1),
        }
    )
    # A flat image has zero variation in it, so std == 0 marks a broken frame.
    # Storing the verdict as a column keeps the rule in one readable place.
    df["is_flat"] = df["std"] == 0
    return df


def normalise(images):
    """Min-max scale to [0, 1] and z-score standardise, both fully vectorised."""
    if not isinstance(images, np.ndarray):
        raise TypeError("expected a NumPy ndarray")

    x = images.astype(np.float32)

    lo, hi = x.min(), x.max()
    span = hi - lo or 1.0  # guard against a batch where every pixel is identical
    minmax = (x - lo) / span

    mean, std = x.mean(), x.std()
    zscore = (x - mean) / (std or 1.0)
    return minmax, zscore, float(mean), float(std)


if __name__ == "__main__":
    raw = generate_dummy_images()
    print(f"raw batch: {raw.shape}  (batch x height x width), dtype={raw.dtype}")

    print("\n=== Pandas: per-image statistics ===")
    stats = describe_batch(raw)
    print(stats.round(2).to_string(index=False))

    print("\n=== Pandas: average brightness per camera ===")
    print(stats.groupby("camera")["mean"].mean().round(2).to_string())

    bad = stats.loc[stats["is_flat"], "image_id"].to_numpy()
    print(f"\nflat (broken) images found: {bad.tolist()}")

    keep = stats.loc[~stats["is_flat"], "image_id"].to_numpy()
    clean = raw[keep]  # NumPy fancy indexing: select rows by an array of positions
    print(f"kept {len(clean)} of {len(raw)} images -> {clean.shape}")

    print("\n=== NumPy: vectorised normalisation ===")
    minmax, zscore, mean, std = normalise(clean)
    print(f"batch mean = {mean:.4f} | batch std = {std:.4f}")
    print(f"after min-max: min = {minmax.min():.4f}, max = {minmax.max():.4f}")
    print("  (expected exactly 0 and 1)")
    print(f"after z-score: mean = {zscore.mean():.4f}, std = {zscore.std():.4f}")
    print("  (expected ~0 and ~1)")

Running it produces:

raw batch: (12, 28, 28)  (batch x height x width), dtype=uint8

=== Pandas: per-image statistics ===
 image_id camera   mean   std  min  max  is_flat
        0  cam-a  39.05 17.63    0   92    False
        1  cam-b  53.98 18.26    0  112    False
        2  cam-a  68.30 18.04   14  122    False
        3  cam-b   0.00  0.00    0    0     True
        4  cam-a  99.43 17.88   50  155    False
        5  cam-b 114.46 17.77   58  177    False
        6  cam-a 129.40 17.90   74  188    False
        7  cam-b 145.08 18.69   65  202    False
        8  cam-a 255.00  0.00  255  255     True
        9  cam-b 173.74 17.43  120  234    False
       10  cam-a 189.84 17.58  126  255    False
       11  cam-b 203.57 18.80  146  255    False

=== Pandas: average brightness per camera ===
camera
cam-a    130.17
cam-b    115.14

flat (broken) images found: [3, 8]
kept 10 of 12 images -> (10, 28, 28)

=== NumPy: vectorised normalisation ===
batch mean = 121.6832 | batch std = 57.1982
after min-max: min = 0.0000, max = 1.0000
  (expected exactly 0 and 1)
after z-score: mean = -0.0000, std = 1.0000
  (expected ~0 and ~1)

Read this output in four steps, following the order the program runs. One: the mean column climbs from 39 to 203 — those are the different brightness levels we deliberately created, proving the Pandas table really does say something about the data. Two: images 3 and 8 have std = 0. A standard deviation of zero means every pixel is identical, i.e. a flat frame carrying no information — that is how you spot a broken image without looking at it. Three: those two rows get dropped, leaving 10 of 12 images. Four: after normalisation, min-max gives exactly 0 and 1, and z-score gives mean 0 and standard deviation 1 — which is the definition, so this doubles as both a result and a self-check.

The -0.0000 on the last line is not an error: it is floating-point negative zero, arising when the accumulated rounding errors lean a vanishing amount negative. It equals zero.

One detail worth noticing in generate_dummy_images: there is one for loop, and it breaks no rule. That loop runs 12 times — once per image — not 9,408 times, once per pixel. This is the practical boundary of the "don't use loops" rule: looping a few dozen times in the Python layer is fine, looping millions of times is the problem. Every pixel-level calculation in this file is vectorised.

How to run this project on your own machine

1. Install the libraries if you haven't: pip install numpy pandas in your virtual environment.
2. Download image_normalize.py at the end of the lesson, or retype the code above.
3. Run it: python3 image_normalize.py.
4. You will get exactly the numbers above, because rng = np.random.default_rng(42) fixes the random seed. Change 42 to something else and the data changes but the shape of the result doesn't — the two flat images are still detected. Try this too: delete the line images[3] = 0, run again, and watch the list of broken images shrink to one entry.

Lesson summary & bridge to what's next

Bridge to the next lesson: NumPy is very powerful on the CPU, but to train complex deep neural networks on a GPU we need to move to the specialised PyTorch library and its automatic differentiation mechanism, Autograd, in Lesson 4.