NumPy Performance Optimization

Real-World Scenario

A computer vision team processes 10,000 high-resolution images per second for an autonomous vehicle system. Their NumPy pipeline was bottlenecking at 2,000 images per second. After profiling, switching dtypes, fixing memory layout, and eliminating unnecessary copies, they hit the 10k target without changing hardware. The same code — just NumPy used correctly.

Profiling First

Never optimize without measuring. Profile with timeit and memory_profiler to find the actual bottleneck.

import numpy as np
import timeit

rng = np.random.default_rng(42)
A = rng.standard_normal((1000, 1000))
B = rng.standard_normal((1000, 1000))

# timeit for micro-benchmarks
t = timeit.timeit(lambda: A @ B, number=100) / 100
print(f"matmul: {t*1000:.2f}ms per call")

# Use %timeit in Jupyter for convenience:
# %timeit A @ B

# Memory usage with tracemalloc
import tracemalloc
tracemalloc.start()
C = A + B
current, peak = tracemalloc.get_traced_memory()
tracemalloc.stop()
print(f"Peak memory: {peak / 1024:.1f} KB")

Choose the Right dtype

Smaller dtypes mean more data fits in CPU cache, and SIMD instructions process more elements per cycle.

import numpy as np

n = 10_000_000

# Memory comparison
f64 = np.ones(n, dtype=np.float64)
f32 = np.ones(n, dtype=np.float32)
f16 = np.ones(n, dtype=np.float16)
i32 = np.ones(n, dtype=np.int32)
i8  = np.ones(n, dtype=np.int8)

print(f"float64: {f64.nbytes / 1e6:.0f} MB")  # 80 MB
print(f"float32: {f32.nbytes / 1e6:.0f} MB")  # 40 MB
print(f"float16: {f16.nbytes / 1e6:.0f} MB")  # 20 MB
print(f"int32:   {i32.nbytes / 1e6:.0f} MB")  # 40 MB
print(f"int8:    {i8.nbytes  / 1e6:.0f} MB")  # 10 MB

# For ML: use float32 to double throughput vs float64
rng = np.random.default_rng(0)
X_f64 = rng.standard_normal((5000, 5000))
X_f32 = X_f64.astype(np.float32)

import timeit
t64 = timeit.timeit(lambda: X_f64 @ X_f64.T, number=5) / 5
t32 = timeit.timeit(lambda: X_f32 @ X_f32.T, number=5) / 5
print(f"float64 matmul: {t64:.3f}s")
print(f"float32 matmul: {t32:.3f}s")
print(f"Speedup: {t64/t32:.1f}x")  # typically 1.5–2x on modern CPUs

Avoid Unnecessary Copies

Every unnecessary copy doubles memory and adds allocation time. Most NumPy operations return views — only explicit .copy(), fancy indexing, and type casts allocate new memory.

import numpy as np

arr = np.arange(1_000_000, dtype=np.float64)

# Check if an operation is a view or copy
view = arr[::2]          # stride slice → view
copy = arr[[0, 2, 4]]    # fancy index → copy

print(view.base is arr)  # True — view
print(copy.base is None) # True — copy (base is None for owned arrays)

# Avoid unnecessary type casts in hot loops
# Bad: creates a temporary float64 copy on every call
def bad_normalize(x: np.ndarray) -> np.ndarray:
    return (x - x.mean()) / x.std()  # produces float64 even if x is float32

# Good: keep computation in the original dtype
def good_normalize(x: np.ndarray) -> np.ndarray:
    mean = x.mean()
    std = x.std()
    out = np.empty_like(x)  # pre-allocate output with same dtype
    np.subtract(x, mean, out=out)   # in-place operation on pre-allocated buffer
    np.divide(out, std, out=out)
    return out

Memory Layout: C-order vs F-order

Access patterns aligned with memory layout get cache hits. Misaligned access causes cache misses that can 10x your runtime.

import numpy as np
import timeit

n = 2000

# C-order (default): row-major — A[0,0], A[0,1], A[0,2], ..., A[1,0]
A_c = np.random.default_rng(0).standard_normal((n, n))

# F-order: column-major — A[0,0], A[1,0], A[2,0], ..., A[0,1]
A_f = np.asfortranarray(A_c)

# Row sum: iterates over rows — fast for C-order (cache-friendly)
t_c = timeit.timeit(lambda: A_c.sum(axis=1), number=100) / 100
t_f = timeit.timeit(lambda: A_f.sum(axis=1), number=100) / 100
print(f"Row sum — C-order: {t_c*1000:.2f}ms, F-order: {t_f*1000:.2f}ms")

# Column sum: iterates over columns — fast for F-order
t_c = timeit.timeit(lambda: A_c.sum(axis=0), number=100) / 100
t_f = timeit.timeit(lambda: A_f.sum(axis=0), number=100) / 100
print(f"Col sum — C-order: {t_c*1000:.2f}ms, F-order: {t_f*1000:.2f}ms")

# Ensure contiguous layout before repeated operations
A_view = A_c[::2]  # non-contiguous stride
A_contiguous = np.ascontiguousarray(A_view)  # forces contiguous copy

In-Place Operations

Avoid temporary array allocation by using out= parameters and in-place operators.

import numpy as np

rng = np.random.default_rng(42)
X = rng.standard_normal((10_000, 100))

# Bad: creates 3 temporary arrays (X**2, sum, sqrt)
norms_bad = np.sqrt((X ** 2).sum(axis=1))

# Good: pre-allocate output buffers, reuse them
squared = np.empty_like(X)
np.multiply(X, X, out=squared)         # in-place square
row_sums = squared.sum(axis=1)         # (10000,)
norms_good = np.sqrt(row_sums, out=row_sums)  # sqrt in-place

print(np.allclose(norms_bad, norms_good))  # True

# In-place arithmetic operators
arr = np.ones(1_000_000)
arr += 1.0   # in-place: no new array allocated
arr *= 2.0   # in-place

Using einsum for Complex Reductions

np.einsum expresses multi-dimensional contractions concisely and can be faster than chaining matmul + sum.

import numpy as np

rng = np.random.default_rng(0)
A = rng.standard_normal((100, 50))
B = rng.standard_normal((50, 80))

# Matrix multiplication: output[i,j] = sum_k A[i,k] * B[k,j]
C = np.einsum('ik,kj->ij', A, B)
print(np.allclose(C, A @ B))  # True

# Batch matrix multiply: (batch, n, k) × (batch, k, m)
batch_A = rng.standard_normal((32, 10, 20))
batch_B = rng.standard_normal((32, 20, 15))
result = np.einsum('bnk,bkm->bnm', batch_A, batch_B)
print(result.shape)  # (32, 10, 15)

# Outer product, trace, element-wise batch operations — all via einsum
trace = np.einsum('ii->', np.eye(5))  # sum of diagonal = 5.0
print(trace)

Numba for JIT Compilation

When a computation genuinely requires a loop (irregular data, custom logic), use Numba to JIT-compile it to native CPU code.

import numpy as np
import numba  # pip install numba

@numba.jit(nopython=True, parallel=True)
def pairwise_distances_numba(X: np.ndarray) -> np.ndarray:
    """Compute all pairwise Euclidean distances — O(n^2) with Numba parallelism."""
    n = X.shape[0]
    out = np.empty((n, n), dtype=np.float64)
    for i in numba.prange(n):   # prange → parallel across iterations
        for j in range(n):
            diff = X[i] - X[j]
            out[i, j] = np.sqrt((diff * diff).sum())
    return out

X = np.random.default_rng(0).standard_normal((500, 10))
D = pairwise_distances_numba(X)   # first call triggers JIT compilation
print(D.shape)  # (500, 500)

Performance Checklist