Desempenho: temporários, in-place e dtypes

O custo dos temporários

Cada operação aritmética cria um array novo para o resultado. Em uma expressão com várias operações, o NumPy cria vários temporários, cada um do tamanho do array original. Para arrays grandes, isso significa memória gasta e páginas novas sendo alocadas. A alternativa é escrever no mesmo buffer, com out= e operadores in-place (+=, -=):

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

O resultado é o mesmo. A diferença está na memória: a versão com buffers reutilizados não aloca arrays novos a cada passo. Dá para medir o pico de memória de cada uma:

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

Esse padrão (alocar os buffers uma vez, fora de um laço, e reutilizá-los) é o que bibliotecas de alto desempenho fazem. Para código que roda uma vez, a expressão simples é mais legível e basta. Eu só uso out= quando a medição mostra que a memória ou as alocações importam.

Escolher o dtype certo

O tamanho do tipo define quanta memória o array ocupa, e quanto precisa ser lido da memória a cada operação. Para a maioria dos problemas de aprendizado de máquina, float32 basta, e usa metade:

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

True
True

O custo é a precisão (capítulo 7). Não troque o tipo sem saber quanta precisão o seu problema exige.

Iterar um array em Python é lento

Percorrer um array com for cria um objeto escalar do NumPy a cada passo, e isso é caro. Quando um laço Python é inevitável, converter para lista antes com .tolist() o deixa mais rápido. E funções Python como o sum sobre um array são muito mais lentas do que o método do próprio array:

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

O que mais existe

A ordem que eu sigo

Primeiro, o algoritmo (a vetorização certa). Depois, a memória (tipos e temporários). Só então ferramentas externas de compilação. E, entre uma etapa e outra, medir (capítulo 18): o gargalo quase nunca está onde a intuição aponta.