Speicherwachstum mit Broadcast-Operationen in NumPy

Question

Ich benutze NumPy, um einige große Datenmatrizen (von etwa ~ 50 GB Größe) zu behandeln. Die Maschine, auf der ich diesen Code ausführe, hat 128 GB RAM, so dass einfache lineare Operationen dieser Größenordnung kein Speicherproblem darstellen sollten.

Allerdings bin ich Zeuge eines großen Speicherwachstum (mehr als 100 GB), wenn Sie den folgenden Code in Python Berechnung:

import numpy as np 

# memory allocations (everything works fine) 
a = np.zeros((1192953, 192, 32), dtype='f8') 
b = np.zeros((1192953, 192), dtype='f8') 
c = np.zeros((192, 32), dtype='f8') 

a[:] = b[:, :, np.newaxis] - c[np.newaxis, :, :] # memory explodes here 

Bitte beachten Sie, dass die anfängliche Speicherzuordnungen erfolgt ohne Probleme. Wenn ich jedoch versuche, den Subtraktionsvorgang mit dem Senden auszuführen, wird der Speicher auf mehr als 100 GB erhöht. Ich dachte immer, dass Broadcasting zusätzliche Speicherzuweisungen vermeiden würde, aber jetzt bin ich mir nicht sicher, ob das immer der Fall ist.

Kann also jemand etwas darüber sagen, warum dieses Speicherwachstum stattfindet und wie der folgende Code mit mehr speichereffizienten Konstrukten umgeschrieben werden könnte?

Ich führe den Code in Python 2.7 in IPython Notebook.

Answer 1

@ rths Vorschlag, die Operation in kleineren Chargen durchzuführen, ist eine gute Idee. Sie können auch versuchen, die Funktion np.subtract zu verwenden und ihr das Ziel-Array zu geben, um zu vermeiden, dass ein zusätzliches temporäres Array erstellt wird. Ich denke auch, dass Sie c als c[np.newaxis, :, :] nicht indexieren müssen, weil es bereits ein 3-d-Array ist.

Also statt

a[:] = b[:, :, np.newaxis] - c[np.newaxis, :, :] # memory explodes here 

versuchen

np.subtract(b[:, :, np.newaxis], c, a) 

Das dritte Argument von np.subtract das Array Ziel.

Answer 2

Nun, Ihr Array a dauert bereits 1192953*192*32* 8 bytes/1.e9 = 58 GB des Speichers.

Die Ausstrahlung macht keine zusätzliche Speicherzuordnungen für die erste Arrays, sondern das Ergebnis des

b[:, :, np.newaxis] - c[np.newaxis, :, :] 

immer noch in einer temporären Matrix gespeichert wird. Daher haben Sie in dieser Zeile mindestens 2 Arrays mit der Form a für den gesamten belegten Speicher reserviert >116 GB.

Sie können dieses Problem vermeiden, indem Sie auf einmal auf einem kleineren Teil Ihrer Array arbeitet,

CHUNK_SIZE = 100000 
for idx in range(b.shape[0]/CHUNK_SIZE): 
    sl = slice(idx*CHUNK_SIZE, (idx+1)*CHUNK_SIZE) 
    a[sl] = b[sl, :, np.newaxis] - c[np.newaxis, :, :] 

dies wird geringfügig langsamer, aber verbraucht viel weniger Speicher.