Write Faster Python for ML: Vectorization, Profiling, and Memory Management

Introduction

Fast ML code in Python comes from three disciplines applied in order: measure where time and memory go, push the heavy work out of the Python interpreter, and control how much data lives in memory at once. Skip the first and you optimize the wrong thing; skip the other two and no amount of hardware will save you.

Python is an unusual language for high-performance work. The interpreter itself is slow: a simple loop that adds numbers runs tens to hundreds of times slower than the same loop in C. Yet Python dominates machine learning, because the libraries that matter (NumPy, SciPy, scikit-learn, PyTorch, JAX) do their real work in compiled C, C++, CUDA or Fortran. Python is the steering wheel, not the engine.

That gives a simple rule for performance: the less time your program spends executing Python bytecode per data element, the faster it runs. Vectorization, compilation and good memory layout are all ways of following that rule.

Why ML code gets slow

In practice, slow ML code usually has one of a few causes:

Python loops over data. Iterating over rows, pixels, tokens or samples one at a time in pure Python.

Excess memory traffic. Creating large temporary arrays, copying data needlessly, or using 64-bit types where 32 or 16 would do. Modern CPUs and GPUs are often limited by memory bandwidth, not arithmetic.

Waiting on I/O. Training loops that stall while data is read and decoded.

Running out of memory. Swapping, crashing, or being forced to use tiny batches.

Optimizing blind. Spending a day speeding up a function that accounts for 3% of runtime.

This article tackles these in the order you should: profiling first, then vectorization and compilation for CPU time, then memory management from RAM to GPU. It closes with a case study that applies everything to one pipeline, and a checklist for code reviews.

Timings quoted in this article are typical orders of magnitude on a modern laptop. Your numbers will differ with hardware, data shape and library versions, which is exactly why the next section comes first.

Profiling CPU time: find the real bottleneck

Profile before you optimize. In most ML codebases, a small fraction of the code accounts for most of the runtime, and it is rarely the part you would guess.

Use tools in a funnel: start coarse to find the slow stage, then zoom in to the slow function, then the slow line.

Micro-benchmarks with timeit

timeit runs a snippet many times and reports a robust average, which avoids the noise of a single measurement.

import numpy as np
rng = np.random.default_rng(0)       # modern Generator API
x = rng.random(1_000_000)

%timeit sum(v * v for v in x)        # Python loop: ~100s of ms
%timeit np.dot(x, x)                 # vectorized:  ~0.5 ms

Outside notebooks, use timeit.repeat and take the minimum of several runs: the minimum is the best estimate of the code's true cost, while higher values reflect interference from other processes.

Function-level profiling with cProfile

cProfile ships with Python and records every function call. Sort by cumulative time to see which top-level steps are expensive, and by total time to see which functions do the work themselves.

python -m cProfile -o train.prof train.py
snakeviz train.prof         # interactive flame-style view in the browser

import cProfile, pstats

with cProfile.Profile() as pr:
    build_features(df)
pstats.Stats(pr).sort_stats("cumulative").print_stats(15)

A limitation to keep in mind: cProfile sees Python functions, so time spent inside a single NumPy call appears as one opaque entry. That is fine; it tells you the time is already in compiled code.

Line-level profiling

Once you know the hot function, line_profiler shows time per line. Decorate the function and run with kernprof. With line_profiler 4.1 or later, import profile explicitly: a bare @profile raises NameError when the module runs outside kernprof, while the imported decorator is a no-op until you enable it with LINE_PROFILE=1.

from line_profiler import profile   # no-op unless profiling is enabled

@profile
def build_features(df):
    df["ratio"] = df["a"] / df["b"]
    df["bucket"] = df.apply(bucketize, axis=1)   # <- usually lights up
    return df

kernprof -l -v build_features.py          # classic: run under kernprof
LINE_PROFILE=1 python build_features.py   # line_profiler 4.1+

Sampling profilers for real workloads

py-spy attaches to a running process and samples its call stack, so you can profile a long training job or a production service without modifying or restarting it. py-spy top --pid <PID> gives a live view; py-spy record -o profile.svg --pid <PID> writes a flame graph. On Python 3.15 and later, the standard library also ships a sampling profiler, profiling.sampling (for example python -m profiling.sampling attach <PID>), which helps where installing third-party tools is not allowed.

Scalene is worth knowing because it splits time into Python versus native code and reports memory per line. If a line shows mostly Python time, it is a vectorization or compilation candidate. If it shows mostly native time, the code is already efficient and gains must come from doing less work or using better hardware.

Profiling hygiene

Profile on realistic data sizes; a 1,000-row sample hides problems that appear at 10 million.

Warm up first. JIT compilers, caches and GPU kernels have a first-call cost that distorts single measurements.

On GPU, call torch.cuda.synchronize() before starting and before stopping a timer (or time with a pair of torch.cuda.Event objects), because GPU work runs asynchronously.

Change one thing at a time, and keep a log of before-and-after timings.

Vectorization fundamentals

Vectorization means expressing an operation on whole arrays so that the loop runs inside compiled code rather than in Python. It is the single highest-leverage technique for CPU-bound ML code, routinely giving 10–1000x speed-ups.

Why it is fast

A NumPy array is a contiguous block of memory holding values of one type. When you write a + b, NumPy runs a tight C loop over that memory with no type checks per element, and often uses SIMD instructions that process several values per CPU cycle. A Python loop, by contrast, unpacks a Python object, checks its type, dispatches the addition and allocates a new object for every single element.

Universal functions

NumPy's universal functions (ufuncs) apply element-wise: arithmetic, np.exp, np.log, np.sqrt, np.maximum, comparisons and many more. Most numeric feature transforms are a few ufuncs chained together.

# Python loop
out = [1 / (1 + math.exp(-v)) for v in z]

# Vectorized sigmoid (in production, scipy.special.expit(z) gives the
# same result without overflow warnings for large negative z)
out = 1 / (1 + np.exp(-z))

# Numerically stable softmax over rows
def softmax(logits):
    shifted = logits - logits.max(axis=1, keepdims=True)
    e = np.exp(shifted)
    return e / e.sum(axis=1, keepdims=True)

Broadcasting

Broadcasting lets arrays of different shapes combine without explicit loops or copies. NumPy compares shapes from the rightmost dimension; two dimensions are compatible if they are equal or one of them is 1. A dimension of size 1 is virtually stretched to match.

rng = np.random.default_rng(0)
X = rng.random((100_000, 32), dtype=np.float32)  # (samples, features)
mu = X.mean(axis=0)                       # (32,)
sigma = X.std(axis=0)                     # (32,)
sigma[sigma == 0] = 1.0                   # guard constant features
X_std = (X - mu) / sigma                  # (100000, 32) - (32,) works

The keepdims=True argument and None (or np.newaxis) indexing are the two tools for lining up shapes deliberately:

row_norms = np.linalg.norm(X, axis=1, keepdims=True)   # (100000, 1)
X_unit = X / np.maximum(row_norms, 1e-12)  # normalize; zero-row safe

# Pairwise differences between two sets of points
A = rng.random((500, 3))
B = rng.random((800, 3))
diff = A[:, None, :] - B[None, :, :]                    # (500, 800, 3)
dist = np.sqrt((diff ** 2).sum(axis=-1))                # (500, 800)

That last example is elegant but materializes a 500 × 800 × 3 intermediate. For large inputs, the algebraic expansion shown in the next section uses far less memory.

Boolean masks instead of if-statements

Conditional logic inside loops becomes masking. Comparisons produce boolean arrays you can use to select, count or assign.

# Clip outliers and flag them
limit = np.percentile(x, 99)
is_outlier = x > limit
x_clipped = np.where(is_outlier, limit, x)
outlier_rate = is_outlier.mean()

# Multi-branch logic
labels = np.select([score >= 0.8, score >= 0.5], ["high", "medium"], default="low")

Thinking in arrays

The skill of vectorization is mostly a change of habit. Instead of asking "what do I do with each element?", ask "what operation transforms this whole array into that one?" A useful exercise is to write the loop version first, then identify the shape of the inputs and outputs, and finally express the loop body as array operations with the loop index removed.

Advanced vectorization patterns

Once the basics are second nature, a handful of patterns cover most of the remaining loops in ML code: matrix formulations, einsum, grouped reductions, sliding windows, and avoiding functions that only look vectorized.

Rewrite as matrix multiplication

Matrix multiplication is the most optimized operation in numerical computing. NumPy dispatches @ to a BLAS library (OpenBLAS, MKL or Accelerate) that is multi-threaded and cache-tuned. If you can express a computation as a matrix product, do it. BLAS is used only for floating-point and complex dtypes; an integer matrix product falls back to a much slower generic loop, so cast to float32 first.

The pairwise-distance example from the previous section becomes:

# ||a - b||^2 = ||a||^2 + ||b||^2 - 2 a.b
sq_a = (A ** 2).sum(axis=1)[:, None]      # (500, 1)
sq_b = (B ** 2).sum(axis=1)[None, :]      # (1, 800)
d2 = sq_a + sq_b - 2 * A @ B.T            # (500, 800), no 3-D temporary
d2 = np.maximum(d2, 0)                    # guard against tiny negative rounding

This is how scikit-learn, FAISS and most nearest-neighbour libraries compute distances, and it is often 10x faster than the broadcasting version while using a fraction of the memory.

einsum for complex contractions

np.einsum expresses sums over products using index notation. It replaces stacks of transposes, reshapes and loops with one readable line. Set optimize=True to let NumPy choose an efficient contraction order.

# Batched matrix multiply: (batch, i, k) x (batch, k, j)
C = np.einsum("bik,bkj->bij", A, B)       # same result as A @ B

# Row-wise dot products of two (n, d) arrays
dots = np.einsum("nd,nd->n", U, V)

# Attention scores: queries (b, q, d) against keys (b, k, d)
d = Q.shape[-1]
scores = np.einsum("bqd,bkd->bqk", Q, K, optimize=True) / np.sqrt(d)

The same notation works in PyTorch (torch.einsum) and JAX, so code ports cleanly. One caveat: without optimize=True, np.einsum runs its own loops rather than BLAS, so for a plain or batched matrix product A @ B (np.matmul) is usually faster. Reach for einsum when the contraction has no direct matmul equivalent.

Grouped operations without Python loops

Per-group aggregation is a common loop: "for each user, sum their values". With integer group IDs, np.bincount and np.add.at do it in one call.

# group_ids: integer labels 0..G-1, values: same length
group_sum = np.bincount(group_ids, weights=values, minlength=G)
group_count = np.bincount(group_ids, minlength=G)
group_mean = group_sum / np.maximum(group_count, 1)

# Broadcast the group mean back to every row
centered = values - group_mean[group_ids]

If IDs are strings or arbitrary integers, map them to compact codes first with np.unique(ids, return_inverse=True) or pd.factorize.

Sliding windows

Rolling features over time series or signals are often written as loops over window start positions. sliding_window_view creates a strided view with no data copy, which you can then reduce along the window axis.

from numpy.lib.stride_tricks import sliding_window_view

signal = np.random.default_rng(0).random(1_000_000)
win = sliding_window_view(signal, window_shape=50)   # (999951, 50) view
rolling_mean = win.mean(axis=1)
rolling_max = win.max(axis=1)

For simple rolling sums or means, a cumulative sum is even faster, because each window becomes a single subtraction:

c = np.concatenate(([0.0], np.cumsum(signal, dtype=np.float64)))
rolling_mean = (c[50:] - c[:-50]) / 50

Accumulate the cumulative sum in float64, even for float32 signals: over millions of elements, rounding error in a low-precision running total grows large enough to distort the window differences.

Beware of fake vectorization

Some functions look vectorized but loop in Python under the hood:

np.vectorize is a convenience wrapper around a Python loop. It gives broadcasting semantics, not speed.

np.apply_along_axis calls a Python function once per slice.

pandas.DataFrame.apply with axis=1 calls a Python function once per row.

List comprehensions building arrays still execute Python bytecode per element.

If profiling shows one of these is hot, the fix is a true array expression or the compilation tools in the next section.

When vectorization isn't enough: compiling Python

Some algorithms are inherently sequential or would need huge temporary arrays to vectorize. For those, compile the loop instead of eliminating it. Numba is the easiest option for NumPy-style code; JAX fits differentiable and accelerator workloads; Cython suits library authors who need fine control.

Typical candidates

Recurrences where each step depends on the previous one: running balances, exponential smoothing with resets, dynamic programming, custom RNN-like updates.

Algorithms with early exits or data-dependent branching, such as custom tree traversal or simulations.

Computations where the vectorized form would create intermediates many times larger than the inputs.

Numba: decorate and go

Numba compiles a subset of Python and NumPy to machine code via LLVM the first time a function is called. Write the plain loop, add @njit, and pass NumPy arrays.

import numpy as np
from numba import njit, prange

@njit(cache=True)
def ewma_with_resets(x, resets, alpha):
    out = np.empty_like(x)
    s = x[0]
    for i in range(x.shape[0]):
        if resets[i]:
            s = x[i]
        else:
            s = alpha * x[i] + (1 - alpha) * s
        out[i] = s
    return out

A loop like this typically runs 50–200x faster than pure Python and close to hand-written C. cache=True saves the compiled code to disk so later runs skip compilation.

For independent iterations, parallel=True with prange spreads the loop across CPU cores:

@njit(parallel=True, cache=True)
def row_l2_norms(X):
    n, d = X.shape
    out = np.empty(n, dtype=X.dtype)
    for i in prange(n):
        acc = 0.0
        for j in range(d):            # scalar loop: no temporaries
            acc += X[i, j] * X[i, j]
        out[i] = np.sqrt(acc)
    return out

Writing the inner loop as scalar code avoids allocating a temporary array for every row, which matters inside prange. For this particular operation, np.sqrt(np.einsum("ij,ij->i", X, X)) is already fast; parallel Numba earns its keep when the per-row work is too irregular to express as array operations.

Numba works best on numeric arrays and simple Python types. It does not accelerate Pandas objects, arbitrary classes or most third-party library calls, so extract .to_numpy() arrays before calling a jitted function.

JAX: vectorize, compile and differentiate

JAX offers a NumPy-like API (jax.numpy) with three composable transformations: jit compiles a function with XLA for CPU, GPU or TPU; vmap automatically vectorizes a function written for a single example over a batch; and grad computes derivatives.

import jax, jax.numpy as jnp

def loss_one(w, x, y):                 # written for ONE example
    pred = jnp.dot(x, w)
    return (pred - y) ** 2

batched_loss = jax.vmap(loss_one, in_axes=(None, 0, 0))

@jax.jit
def mean_loss(w, X, Y):
    return batched_loss(w, X, Y).mean()

grad_fn = jax.jit(jax.grad(mean_loss))

vmap is particularly useful: it lets you write clear per-example code and still get fully vectorized execution. JAX requires pure functions and immutable arrays, which is a mindset shift from NumPy.

Cython and friends

Cython compiles annotated Python to C extensions. It gives the most control, including releasing the GIL and calling C libraries directly, but it adds a build step. It is the right tool when you are building a reusable library, less so for one-off analysis code. Many core parts of scikit-learn and Pandas are written in Cython.

Choosing

A note on threads: CPython's global interpreter lock (GIL) has historically prevented Python threads from running bytecode in parallel. NumPy, Numba (nogil=True) and most compiled libraries release it during heavy computation, so threads still help there. Free-threaded CPython was introduced as experimental in 3.13 and became officially supported in 3.14 (PEP 779), but it is still an optional, separate build (for example python3.14t): the default interpreter keeps the GIL, and not every compiled dependency ships free-threaded wheels yet. For production pipelines today, multiprocessing or compiled parallel loops remain the reliable route for pure-Python parallelism; free threading is worth benchmarking once your whole dependency stack supports it.

Memory fundamentals

Memory management affects speed as much as capacity. Smaller arrays fit in CPU caches, move faster across the memory bus, and leave room for bigger batches. Four levers matter most: data types, views versus copies, temporary arrays, and memory layout.

Choose the smallest safe dtype

NumPy defaults to float64 and int64. Most ML workloads do not need that precision.

A 10 million × 100 feature matrix is 8 GB as float64 and 4 GB as float32. Converting once, early, halves memory and usually speeds up arithmetic too.

X = np.load("features.npy").astype(np.float32, copy=False)
print(X.nbytes / 1e9, "GB")
copy=False only skips the copy when the data is already float32. If the file holds float64, this line briefly keeps both versions in memory (12 GB at peak for the example above). In production, convert once and save the float32 result, or convert chunk by chunk from a memory-mapped file (see Working with data larger than RAM).

Watch for silent upcasting: combining a float32 array with a float64 array or a NumPy float64 scalar produces float64. Under NumPy 2's promotion rules (NEP 50), a plain Python float such as 0.5 keeps the array's dtype, but an np.float64 value, such as a statistic computed from a float64 array, does not. Keep constants in the same dtype (np.float32(0.5)) in hot code, and assert X.dtype at stage boundaries.

Views versus copies

Basic slicing returns a view: a new array object that shares the original memory. Fancy indexing (integer arrays or boolean masks) returns a copy.

X = np.zeros((1_000_000, 50), dtype=np.float32)

first_half = X[:500_000]          # view, no new memory
every_other = X[::2]              # view (strided)
selected = X[[0, 10, 20]]         # copy
positives = X[X[:, 0] > 0]        # copy

np.shares_memory(X, first_half)   # True

Views are free, but they keep the entire parent array alive. If you slice a small piece from a huge array and discard the original, the full allocation stays in memory until the view is gone. Call .copy() on the small piece to release the parent.

reshape, transpose and .T usually return views too. flatten always copies; ravel and np.ascontiguousarray copy only when the input is not already C-contiguous.

Avoid needless temporaries

Each operation in an expression like y = a * b + c * d - e allocates a full-size temporary. For large arrays, temporaries can double or triple peak memory. In-place operators and the out= argument reuse existing buffers.

# Allocates several temporaries
y = np.exp(-(x - mu) ** 2 / (2 * sigma ** 2))

# Reuses one buffer
y = x - mu
np.square(y, out=y)
y /= -(2 * sigma ** 2)
np.exp(y, out=y)

The second form is less readable, so reserve it for hot paths on large arrays that profiling has flagged. The numexpr library offers a middle ground: it evaluates a whole expression string in cache-sized chunks without full temporaries, and uses multiple threads.

Memory layout and contiguity

NumPy arrays are row-major (C order) by default: elements of a row sit next to each other. Iterating or reducing along the last axis follows memory order and is cache-friendly; striding down columns of a huge C-ordered array is slower. After transposes or strided slicing, an array may be non-contiguous, and some libraries silently copy it before use. If you pass the same array to a library many times, make it contiguous once with np.ascontiguousarray.

Python objects are expensive

A Python float is an object of roughly 24 bytes plus an 8-byte pointer in a list, versus 8 bytes (or 4) in a NumPy array. A list of a million floats therefore uses about four times the memory of a float64 array. Lists of dictionaries, a common shape for JSON records, are far worse. Convert to arrays or columnar structures as early as possible.

Profiling memory

Out-of-memory crashes are caused by peak memory, not average memory, and the peak is often a brief moment when several large temporaries coexist. Memory profilers show you exactly when and where that peak happens.

Quick checks

Before reaching for a profiler, a few one-liners answer most questions:

X.nbytes / 1e6                              # MB held by a NumPy array
df.memory_usage(deep=True).sum() / 1e6      # MB held by a DataFrame, incl. strings

import psutil
psutil.Process().memory_info().rss / 1e6    # process resident memory, MB

Logging resident memory (RSS) at the start and end of each pipeline stage, next to step timings, gives a cheap and surprisingly effective memory profile.

tracemalloc: built in

tracemalloc in the standard library tracks allocations made through Python's allocator, including NumPy array buffers. Its most useful feature is reporting the peak.

import tracemalloc

tracemalloc.start()
try:
    features = build_features(raw)
    current, peak = tracemalloc.get_traced_memory()
    print(f"current={current/1e6:.0f} MB  peak={peak/1e6:.0f} MB")

    snapshot = tracemalloc.take_snapshot()
    for stat in snapshot.statistics("lineno")[:10]:
        print(stat)
finally:
    tracemalloc.stop()

If peak is several times larger than both the input and the output, you are creating large temporaries: look for chained arithmetic, intermediate DataFrames, or fancy-indexing copies. Note that the snapshot lists memory still allocated when it is taken, not what existed at the peak; to attribute the peak itself, use memray (below).

memray: detailed and production friendly

Bloomberg's memray tracks every allocation, including those made by C extensions, with low overhead. It produces flame graphs of where memory was allocated at the peak, and it can attach to running processes.

memray run -o train.bin train.py
memray flamegraph train.bin          # HTML flame graph of peak memory
memray summary train.bin             # top allocators in the terminal

There is also a pytest plugin (pytest-memray) that fails tests when a function exceeds a memory limit, which is a good way to stop memory regressions from creeping into a pipeline.

Line-by-line memory

Scalene, mentioned in the profiling section, reports memory growth per line alongside CPU time, making it easy to spot the single line that allocates gigabytes. The older memory_profiler package offers a similar @profile decorator, although it is no longer actively maintained.

Common findings and fixes

One caution on the last row: Python and the C allocator often keep freed memory reserved for reuse instead of returning it to the operating system. RSS that does not drop after del is not necessarily a leak. A true leak shows as growth that continues across repeated iterations of the same work.

Working with data larger than RAM

When a dataset does not fit in memory, stop trying to load it all. Stream it, map it from disk, or process it in chunks, so that memory use depends on the batch size rather than the dataset size.

Generators and streaming

A generator yields items one at a time instead of building a full list. It is the simplest way to keep memory flat when reading files, records or batches.

def batches(paths, batch_size=4096):
    buf = []
    for path in paths:
        with open(path, encoding="utf-8") as f:
            for line in f:
                buf.append(parse(line))
                if len(buf) == batch_size:
                    yield np.asarray(buf, dtype=np.float32)
                    buf = []
    if buf:
        yield np.asarray(buf, dtype=np.float32)

for X_batch in batches(files):
    model.partial_fit(X_batch, ...)

Scikit-learn estimators with partial_fit (such as SGDClassifier, MiniBatchKMeans and IncrementalPCA) are designed for exactly this out-of-core pattern. For classifiers, pass the complete label set on the first call (partial_fit(X, y, classes=all_classes)), because a single batch may not contain every class.

Memory-mapped arrays

A memory-mapped array looks like a normal NumPy array but lives on disk. The operating system pages in only the parts you touch, so you can index into a 200 GB file on a 16 GB machine.

# Write once
X = np.lib.format.open_memmap("features.npy", mode="w+",
                              dtype=np.float32, shape=(50_000_000, 128))
for start, chunk in produce_chunks():
    X[start:start + len(chunk)] = chunk
X.flush()
del X                                    # close the mapping

# Read many times, lazily
X = np.load("features.npy", mmap_mode="r")
batch = X[np.sort(idx_batch)]            # only these rows are read

Memory mapping works best with sequential or batched access on fast local SSDs. Random access to individual rows on network storage will be slow. Sorting each batch's row indices before indexing, as above, turns scattered reads into near-sequential ones without changing which rows are in the batch.

Chunked computation

Many statistics can be computed in a single pass by combining per-chunk results: sums, counts, minimums, maximums, and means and variances via running totals.

n = 0
s = np.zeros(128, dtype=np.float64)
ss = np.zeros(128, dtype=np.float64)
for chunk in chunks_of(X, 1_000_000):
    n += len(chunk)
    s += chunk.sum(axis=0, dtype=np.float64)
    ss += np.square(chunk, dtype=np.float64).sum(axis=0)
mean = s / n
var = np.maximum(ss / n - mean ** 2, 0.0)   # clamp rounding negatives
std = np.sqrt(var)

Note the float64 accumulators: summing millions of float32 values loses precision, so accumulate in higher precision even when data is stored compactly. (When a feature's mean is large relative to its spread, the sum-of-squares formula suffers catastrophic cancellation; Welford's algorithm, or Chan's method for merging per-chunk means and variances, is more numerically stable.)

Columnar formats and Arrow

For tabular data, Parquet files read through PyArrow let you load only selected columns and row groups, and Arrow's in-memory format can be shared between Pandas, Polars, DuckDB and PyTorch with little or no copying. Arrow datasets can also be scanned in record batches, which plugs neatly into the generator pattern above.

import numpy as np
import pyarrow.dataset as ds

dataset = ds.dataset("s3://bucket/features/", format="parquet")
for batch in dataset.to_batches(columns=FEATURES, batch_size=65_536):
    X = np.empty((batch.num_rows, len(FEATURES)), dtype=np.float32)
    for j, name in enumerate(FEATURES):
        col = batch.column(name).fill_null(0.0)    # handle nulls first
        X[:, j] = col.to_numpy(zero_copy_only=False)
    train_step(X)

Two production details matter here. Array.to_numpy() defaults to zero-copy-only mode and raises on nulls or non-primitive types, so handle nulls explicitly (fill them, or keep a mask) and pass zero_copy_only=False. And filling a preallocated float32 array avoids the extra full-size copies that np.column_stack(...).astype(np.float32) creates.

Process boundaries as memory boundaries

If a heavy preprocessing stage leaves memory fragmented, run it in a separate process (for example with multiprocessing or as its own pipeline task) and pass results through files. When the process exits, the operating system reclaims everything. This is a blunt but reliable tool for long-running pipelines.

GPU memory and throughput in deep learning

On GPUs the same principles apply with higher stakes: memory is scarcer, transfers between CPU and GPU are expensive, and an idle GPU wastes money. The goals are to keep the GPU busy, keep data on the device, and fit the largest useful batch.

Where GPU memory goes

During training, GPU memory holds four main things: model parameters, gradients, optimizer state (Adam keeps two extra values per parameter), and activations saved for the backward pass. For large models and batches, activations usually dominate.

Mixed precision

Training in bfloat16 or float16 for most operations roughly halves activation memory and uses the GPU's tensor cores, often giving 1.5–3x higher throughput, usually with no measurable loss in accuracy (confirm on your validation metrics).

import torch

device = torch.device("cuda")
scaler = torch.amp.GradScaler("cuda")   # float16 only; not for bfloat16
for x, y in loader:
    x = x.to(device, non_blocking=True)  # async with pin_memory=True
    y = y.to(device, non_blocking=True)
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        loss = loss_fn(model(x), y)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)           # unscale before clipping
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)
    scaler.update()

On hardware that supports it, bfloat16 is usually simpler because it has the same exponent range as float32 and does not need a gradient scaler.

Fitting bigger effective batches

Gradient accumulation. Run several small micro-batches and step the optimizer once, simulating a larger batch without the memory. Divide each micro-batch loss by the number of accumulation steps so the gradient matches the large-batch average.

Activation checkpointing. torch.utils.checkpoint discards some activations in the forward pass and recomputes them during backward, trading compute for memory.

Leaner optimizers. 8-bit optimizer states or optimizers with smaller state can cut memory for large models.

Sharding. For multi-GPU training, FSDP or DeepSpeed ZeRO split parameters, gradients and optimizer state across devices.

Avoid silent memory leaks and sync points

A classic bug keeps tensors, and their whole computation graph, alive across steps:

# Bug: keeps every step's graph in memory
total_loss += loss

# Fix: detach so the graph is freed; the value stays on the GPU (no sync)
total_loss += loss.detach()

# Convert to a Python number only when logging
if step % log_every == 0:
    print(f"loss={total_loss.item() / log_every:.4f}")
    total_loss = 0.0

Avoid the tempting loss.item() fix inside the hot loop: .item(), .cpu(), .numpy() and print(tensor) force the CPU to wait for the GPU to finish. Calling them every step stalls the pipeline, so accumulate detached tensors on the device, as above, and synchronize only every N steps or at logging time.

Wrap evaluation and inference in torch.no_grad() (or torch.inference_mode()) so that no graph is built at all.

Keep the GPU fed

If GPU utilization in nvidia-smi hovers well below 100%, the bottleneck is usually data loading on the CPU, not the model. Bear in mind that nvidia-smi's utilization figure only measures the fraction of time at least one kernel was running, not how much of the GPU each kernel used, so treat it as a first signal and confirm with the PyTorch Profiler.

loader = torch.utils.data.DataLoader(
    dataset,
    batch_size=256,
    shuffle=True,
    num_workers=8,            # parallel decoding/augmentation
    pin_memory=True,          # faster async host-to-device copies
    persistent_workers=True,  # keep workers alive (num_workers > 0)
    prefetch_factor=4,        # batches per worker (num_workers > 0)
    drop_last=True,           # fixed batch shape helps torch.compile
)

Do expensive, deterministic preprocessing once and cache the result (for example as memory-mapped arrays or tokenized shards) instead of repeating it every epoch. Move augmentations to the GPU when the CPU cannot keep up.

Compile the model

In PyTorch 2.x, torch.compile(model) traces the model and fuses many small operations into fewer, larger GPU kernels. Fusion reduces memory traffic, which is often the true limit, and commonly yields meaningful speed-ups for little effort. Measure with and without it, since gains vary by model. Expect the first iterations to be slow while kernels compile, and keep input shapes stable (or pass dynamic=True) to avoid repeated recompilation.

Inspect memory directly

print(torch.cuda.memory_allocated() / 1e9, "GB allocated by tensors")
print(torch.cuda.memory_reserved() / 1e9, "GB reserved by allocator")
print(torch.cuda.max_memory_allocated() / 1e9, "GB peak")
print(torch.cuda.memory_summary())
torch.cuda.reset_peak_memory_stats()   # measure next stage's peak

PyTorch's caching allocator keeps freed blocks reserved, so nvidia-smi shows more memory in use than tensors actually occupy (the CUDA context itself also takes several hundred MB). Use the functions above for the real picture. For detailed debugging, call torch.cuda.memory._record_memory_history(), run a few steps, save a snapshot with torch.cuda.memory._dump_snapshot("snapshot.pickle"), and open it in PyTorch's memory visualizer (pytorch.org/memory_viz). Both are underscore-prefixed APIs, so pin the PyTorch version in scripts that rely on them.

Case study: speeding up a feature pipeline

This illustrative example applies the article's workflow to a typical feature-building step. The numbers are representative of what this kind of refactor achieves, not a benchmark; the order of the steps is the point.

The starting point

A recommendation team computes, for each of 2 million users, the cosine similarity between the user's embedding and 5,000 item embeddings, keeps the top 20 items, and computes an exponentially weighted engagement score from each user's event history. The nightly job takes about 3 hours and peaks at roughly 60 GB of RAM.

# Original (simplified)
results = []
for uid in user_ids:
    u = user_emb[uid]                                   # float64, (256,)
    sims = [np.dot(u, item_emb[j]) /
            (np.linalg.norm(u) * np.linalg.norm(item_emb[j]))
            for j in range(n_items)]
    top = np.argsort(sims)[::-1][:20]
    score = engagement_score(events[events.user_id == uid])   # Python loop inside
    results.append({"user_id": uid, "top": top, "score": score})
out = pd.DataFrame(results)

Step 1: Profile

A py-spy flame graph and step timers show three hot spots: the similarity list comprehension (about 70% of time), the per-user DataFrame filter events[events.user_id == uid] (about 20%), and engagement_score, a Python loop over each user's events (about 8%). tracemalloc shows the peak comes from loading all event columns as float64 and object.

Step 2: Vectorize similarity as a matrix product

Normalize both embedding tables once, then compute similarities for a block of users with one matrix multiply. Use argpartition to find the top 20 without fully sorting 5,000 scores.

def l2_normalize(E):
    E = E.astype(np.float32)                       # new float32 copy
    norms = np.linalg.norm(E, axis=1, keepdims=True)
    E /= np.maximum(norms, 1e-12)                  # zero-vector guard
    return E

user_n = l2_normalize(user_emb)
item_n = l2_normalize(item_emb)

def top_k_block(U, k=20):
    sims = U @ item_n.T                            # (block, 5000)
    np.negative(sims, out=sims)                    # in place, no copy
    idx = np.argpartition(sims, k, axis=1)[:, :k]  # unordered top-k
    order = np.argsort(np.take_along_axis(sims, idx, axis=1), axis=1)
    return np.take_along_axis(idx, order, axis=1)

Step 3: Process in blocks to bound memory

A full 2,000,000 × 5,000 similarity matrix would need 40 GB in float32. Processing 50,000 users at a time caps each similarity block at 1 GB; the int64 index array that argpartition returns adds about 2 GB more, so size the block to the machine. Results are written into a preallocated array.

BLOCK = 50_000
top_items = np.empty((n_users, 20), dtype=np.int32)
for start in range(0, n_users, BLOCK):
    stop = min(start + BLOCK, n_users)
    top_items[start:stop] = top_k_block(user_n[start:stop])

Step 4: Replace per-user filtering with sort-and-split, and compile the recurrence

Instead of filtering the events table once per user (2 million full scans), sort events by user and time once, find each user's boundaries, and run the engagement recurrence in a Numba-compiled loop over the whole array.

events = events.sort_values(["user_id", "ts"])
uid = events["user_id"].to_numpy()
w = events["weight"].to_numpy(dtype=np.float32)
starts = np.flatnonzero(np.r_[True, uid[1:] != uid[:-1]])
ends = np.r_[starts[1:], len(uid)]

@njit(parallel=True, cache=True)
def engagement(w, starts, ends, alpha):
    out = np.empty(len(starts), dtype=np.float32)
    for g in prange(len(starts)):
        s = 0.0
        for i in range(starts[g], ends[g]):
            s = alpha * w[i] + (1 - alpha) * s
        out[g] = s
    return out

scores = engagement(w, starts, ends, 0.3)
score_user_ids = uid[starts]   # scores[i] is for score_user_ids[i]

Note that scores is ordered by sorted user ID and covers only users who have events, so join it back on score_user_ids (giving users with no events a default score) rather than assuming it lines up with user_ids. Also guard the empty-table case before computing starts.

Step 5: Shrink the inputs

Load only the three event columns needed, as int32, float32 and datetime64, from Parquet. Store embeddings as float32 memory-mapped .npy files so they load instantly and share pages across worker processes.

Result

The job drops from hours to a few minutes and fits on a much smaller machine. No single trick did it: profiling pointed at the right places, vectorization removed the biggest loop, compilation handled the sequential part, and memory discipline made the vectorized version feasible.

Checklist and conclusion

Use this checklist when writing or reviewing performance-sensitive ML code.

Profiling

☐ Stage timings and peak memory are logged on every run.

☐ The hot spot was identified with a profiler before optimizing.

☐ Benchmarks use realistic data sizes, warm-up runs and GPU synchronization.

Vectorization and compilation

☐ No Python loops over samples, rows or elements in hot paths.

☐ Pairwise and batched computations use matrix multiplication or einsum.

☐ Grouped aggregations use bincount, sort-and-split or library group-bys, not per-group filtering.

☐ No np.vectorize, apply_along_axis or row-wise apply in hot paths.

☐ Inherently sequential loops are compiled with Numba (or JAX / Cython).

Memory

☐ Features use float32 (or smaller) unless precision requires more; accumulators use float64.

☐ Only required columns are loaded, with compact dtypes declared at read time.

☐ Large temporaries are avoided with in-place ops, out= or block processing.

☐ Small slices of large arrays are copied before the parent is discarded.

☐ Data larger than RAM is streamed, chunked or memory-mapped.

Deep learning

☐ Mixed precision is enabled where hardware supports it.

☐ The data loader keeps the GPU above roughly 90% utilization.

☐ No per-step .item(), .cpu() or printing of tensors in the training loop.

☐ Evaluation runs under torch.no_grad() or inference_mode().

Conclusion

Writing fast Python for ML is less about clever tricks and more about a repeatable process. Measure first, so effort goes where it matters. Move the heavy lifting out of the interpreter, by vectorizing with NumPy and matrix operations wherever possible and compiling the loops that remain. Treat memory as a first-class resource: compact types, few copies, bounded batches and streaming when data outgrows RAM. Apply the same discipline on the GPU, where keeping the device fed and memory lean pays off directly in training time and cost.

Teams that build these habits into code reviews and shared utilities find that their models train faster, their pipelines run on smaller machines, and their experiments turn around in minutes instead of hours.

Back to all blogs