Performance: temporaries, in-place and dtypes

The cost of temporaries

Each arithmetic operation creates an array , copying everything that already existed. In a loop, this costs time proportional to the square of the size. The right pattern is to accumulate in a list (which grows cheaply) and convert for the result. In an expression with multiple operations, NumPy creates several temporaries, each the size of the original array. For large arrays, this means memory wasted and new pages being allocated. The alternative is to write in the same buffer, with out= and in-place operators (+=, -=):

import time
import tracemalloc

import numpy as np

rng = np.random.default_rng(0)
a = rng.random(2_000_000)
b = rng.random(2_000_000)


def com_temporarios(a, b):
    return (a * 2 + b * 3 - 1) ** 2


def sem_temporarios(a, b, t1, t2):
    np.multiply(a, 2, out=t1)
    np.multiply(b, 3, out=t2)
    t1 += t2
    t1 -= 1
    np.square(t1, out=t1)
    return t1


t1 = np.empty_like(a)
t2 = np.empty_like(a)
print(np.allclose(com_temporarios(a, b), sem_temporarios(a, b, t1, t2)))

True

The result is the same. The difference is in memory: the version with reused buffers does not allocate new arrays at each step. It is possible to measure the peak memory of each one:

def pico_de_memoria(funcao):
    tracemalloc.start()
    funcao()
    _, pico = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    return pico


pico_com = pico_de_memoria(lambda: com_temporarios(a, b))
pico_sem = pico_de_memoria(lambda: sem_temporarios(a, b, t1, t2))
print("o pico é menor sem temporários:", pico_sem < pico_com)
print("e muito menor, menos de 10% do outro:", pico_sem < 0.1 * pico_com)

o pico é menor sem temporários: True
e muito menor, menos de 10% do outro: True

This pattern (allocating the buffers at the end:, outside of a loop, and reusing them) is what high-performance libraries do. For code that runs once, the simple expression is more readable and sufficient. I only use out= when measurement shows that memory or allocations matter.

Choosing the right dtype

The size of the type defines how much memory the array occupies, and how much needs to be read from memory at each operation. For most machine learning problems, float32 is enough, and uses half:

a32 = a.astype(np.float32)
print(a32.nbytes == a.nbytes // 2)
print(np.allclose(a32, a, atol=1e-6))

True
True

The cost is precision (chapter 7). Do not change the type without knowing how much precision your problem requires.

Iterating an array in Python is slow

Iterating over an array with for creates a NumPy scalar object at each step, and that is expensive. When a Python loop is unavoidable, convert to a list beforehand with .tolist() makes it faster. And Python functions like the sum on an array are much slower than the array's own method:

def medir(funcao, repeticoes=3):
    melhor = float("inf")
    for _ in range(repeticoes):
        inicio = time.perf_counter()
        funcao()
        melhor = min(melhor, time.perf_counter() - inicio)
    return melhor


x = np.arange(100_000)
print("iterar a lista é mais rápido:", medir(lambda: [v for v in x.tolist()]) < medir(lambda: [v for v in x]))
print("sum do Python é pelo menos 10 vezes mais lento:", medir(lambda: sum(x)) > 10 * medir(lambda: x.sum()))

iterar a lista é mais rápido: True
sum do Python é pelo menos 10 vezes mais lento: True

What else exists

The order I follow

First, the algorithm (the right vectorization). Then, memory (types and temporaries). Only then external compilation tools. And, between one step and another, measure (chapter 18): the bottleneck is almost never where intuition points.