Dati che non stanno in RAM: memmap, HDF5 e il DataLoader che non esplode

C’è un momento preciso in cui un progetto di machine learning smette di essere un esercizio: quando i dati non stanno più nella memoria del computer. Fino a quel punto pd.read_csv e una lista di array bastano per tutto; da lì in poi il notebook muore a metà della seconda epoca, con un messaggio che parla di memoria e non dice cosa fare.

La soluzione non è comprare RAM. È smettere di portarsi tutti i libri a casa e cominciare a prenderne uno alla volta dallo scaffale — il che richiede che lo scaffale sia ordinato in modo da poter prendere un libro solo senza tirare giù tutta la fila.

La regola sotto tutto il resto

Il posto dove si vince o si perde è la divisione dei compiti dentro il Dataset: il costruttore apre, il metodo di accesso legge. Caricare in __init__ è ciò che fa esplodere la memoria; riaprire il file a ogni chiamata di __getitem__ è ciò che fa morire di fame la GPU. In mezzo c’è tutto il mestiere.

Le tre strade in ordine di complessità crescente. La riga in basso è la ragione per cui non basta scegliere: bisogna misurare cosa succede dopo.

Il modo più semplice: l’array mappato

Se i dati sono un array numerico di forma fissa — spettrogrammi tutti uguali, immagini già ridimensionate, una matrice di feature — non serve nessuna libreria. np.memmap presenta un file su disco come se fosse un array: le pagine vengono lette dal sistema operativo solo quando le si tocca davvero.

import numpy as np
from torch.utils.data import Dataset

# scritto una volta sola, in fase di preparazione
dati = np.lib.format.open_memmap(
    "feature.npy", mode="w+", dtype=np.float32, shape=(2_000_000, 128))
for i, riga in enumerate(sorgente):
    dati[i] = riga
dati.flush()

class DaMemmap(Dataset):
    def __init__(self, percorso, etichette):
        self.x = np.load(percorso, mmap_mode="r")   # NON legge: mappa
        self.y = etichette

    def __len__(self):
        return len(self.x)

    def __getitem__(self, i):
        return torch.from_numpy(np.array(self.x[i])), self.y[i]   # array, non asarray

Un file da otto gigabyte si «apre» istantaneamente e occupa memoria residente solo per le righe realmente lette. Il np.array nel metodo di accesso è la riga che sembra ridondante e non lo è, e va scritta esattamente così: np.asarray — che è la forma che viene in mente per prima — non copia. Restituisce una vista sul file mappato, in sola lettura, e PyTorch avverte che non supporta tensori non scrivibili e che scriverci dentro ha comportamento indefinito. np.array copia la singola riga e il problema sparisce: una lettera di differenza, un avviso che quasi nessuno legge.

Limiti onesti: serve una forma fissa, non c’è compressione, e su un disco di rete le prestazioni crollano. Ma per la maggior parte dei casi in cui qualcuno arriva a cercare «dataset più grande della RAM», questo basta e finisce lì.

Quando serve HDF5: la forma del blocco

HDF5 serve quando in un file solo devono stare cose eterogenee — array di forme diverse, metadati, più insiemi — o quando serve compressione. Il concetto che decide le prestazioni è il chunking: il dataset viene diviso in blocchi di dimensione regolare, indicizzati con un B-tree, ed è il blocco l’unità minima che viene letta dal disco.

import h5py

with h5py.File("dati.h5", "w") as f:
    f.create_dataset(
        "x", shape=(2_000_000, 128), dtype="f4",
        chunks=(512, 128),        # un blocco = 512 righe intere
        compression="gzip", compression_opts=4,
    )

La documentazione di h5py suggerisce blocchi fra 10 KiB e 1 MiB, più grandi per dataset più grandi: la fascia bassa va bene per file su un filesystem, mentre per file su object storage in cloud conviene salire, fino a 2-4 MiB. Ma la dimensione è la metà del problema: l’altra è la forma. «La forma del blocco determina quanti blocchi servono a soddisfare una lettura», e un blocco di (512, 128) è giusto se si leggono righe intere, disastroso se si legge una colonna sola — in quel caso ogni colonna costringe a decomprimere tutte le righe di ogni blocco che la contiene.

Le letture a fette sono efficienti perché vengono tradotte direttamente in selezioni native del formato, quindi f["x"][i] legge davvero solo quel che serve. È il motivo per cui HDF5 sta bene dentro un Dataset — a patto di superare la trappola che viene ora.

La trappola del fork

Aprire il file HDF5 nel costruttore e usarlo dai worker sembra la cosa giusta e produce errori incomprensibili o dati silenziosamente sbagliati. La documentazione di h5py è esplicita: la lettura in parallelo da processi separati funziona senza accorgimenti particolari, ma «è consigliabile aprire il file in modo indipendente in ciascun processo lettore; aprirlo una volta e poi fare fork può causare problemi».

Tradotto in codice: il file si apre pigramente, al primo accesso, che avviene già dentro il worker.

class DaHDF5(Dataset):
    def __init__(self, percorso):
        self.percorso = percorso
        self.f = None                      # niente handle prima del fork
        with h5py.File(percorso, "r") as f:
            self.n = len(f["x"])           # solo la lunghezza, e si richiude

    def __len__(self):
        return self.n

    def __getitem__(self, i):
        if self.f is None:                 # ogni worker apre il suo
            self.f = h5py.File(self.percorso, "r")
        return torch.from_numpy(self.f["x"][i]), self.f["y"][i]

I worker che ingrassano

La seconda trappola è documentata da PyTorch e coglie tutti almeno una volta, perché si manifesta come una perdita di memoria che cresce lentamente e non si spiega. Quando si usano più processi di caricamento, «dopo alcune iterazioni i processi worker consumeranno la stessa quantità di memoria del processo padre per tutti gli oggetti Python del padre a cui i worker accedono».

La causa è che il conteggio dei riferimenti di Python scrive dentro ogni oggetto che tocca, e quella scrittura fa copiare la pagina che la condivisione dopo il fork avrebbe dovuto risparmiare. Una lista di due milioni di stringhe con i percorsi dei file, condivisa fra otto worker, diventa nove copie. La cura, indicata dalla documentazione stessa, è «sostituire gli oggetti Python con rappresentazioni non soggette al conteggio dei riferimenti, come oggetti Pandas, Numpy o PyArrow».

# invece di: self.percorsi = ["img/0001.jpg", ...]      # lista Python
self.percorsi = np.array(["img/0001.jpg", ...])          # un array solo

Quando non c’è nemmeno un indice

Resta il caso in cui i dati non si possono indicizzare per posizione: arrivano da un database, da un server remoto o da log generati mentre l’addestramento gira. È lo scenario per cui esiste IterableDataset, adatto — dice la documentazione — quando «le letture casuali sono costose o addirittura improbabili».

Il prezzo va conosciuto in anticipo: non c’è mescolamento vero. Si può solo tenere un serbatoio di qualche migliaio di esempi e pescare da lì, il che rimescola localmente ma non rompe l’ordine globale. Se il flusso arriva ordinato per classe, il modello vedrà mezz’ora di una classe sola, e nessun serbatoio ragionevole lo salva.

Misurare, non indovinare

Tutte e tre le strade spostano il collo di bottiglia dalla memoria al disco, e conviene verificare dove si è fermato. Il controllo più economico dura due minuti: si esegue un giro del caricatore senza il modello, e si confronta con un giro completo.

import time

t = time.perf_counter()
for lotto in loader:
    pass
print(f"solo dati: {time.perf_counter() - t:.1f} s")

Se il solo caricamento impiega quasi quanto l’addestramento intero, la GPU sta aspettando e il tempo va speso lì: più worker, pin_memory=True — che mette i tensori in memoria bloccata e velocizza il trasferimento verso la GPU — o un formato diverso. Se invece il caricamento è già trascurabile, ottimizzarlo non renderà nulla, per quanto sia la parte più divertente da ottimizzare.

Un dataset più grande della RAM non è un problema di memoria: è un problema di ordine. Il formato giusto è quello la cui unità di lettura coincide con il modo in cui leggerai.

Cosa portarti a casa

Prova prima memmap: nella maggior parte dei casi risolve, in dieci righe e senza dipendenze. Passa a HDF5 quando servono forme eterogenee, metadati o compressione, e allora dimensiona i blocchi sul modo in cui leggerai davvero. Tieni IterableDataset per quando l’indice proprio non esiste, sapendo che rinunci al mescolamento. E in tutti e tre i casi metti in conto le due trappole: il file aperto prima del fork e gli oggetti Python che ingrassano i worker — sono documentate entrambe, e sono entrambe guasti silenziosi, che è il motivo per cui costano tanto.