Tensorflow Speicherverwaltung - Chunking?

Question

Ich habe ein Problem mit der Berechnung einer Reduzierung einer Funktion eines Netzwerks, das durch eine große (200000x200000) Matrix dargestellt wird, die als Abstandsmatrix zwischen Paaren von Punkten erzeugt wird.

Minimal Beispiel Eingabe X einen 200000x2 numpy Array von kartesischen Koordinaten:

x = tf.constant(X[:,0], shape=[X.shape[0],1]) 
y = tf.constant(X[:,1], shape=[X.shape[0],1]) 
dx = x - tf.transpose(x) 
dy = y - tf.transpose(y) 
D = tf.sqrt(dx*dx + dy*dy) 
M = 0.1 * 5.0/tf.pow(4.0 + D, 1.5) 
res = tf.reduce_sum(betaM) 

auf der CPU läuft, den Speicher (16 GB auf meinem MBP) ist schnell überzeichnet und das System zum Stillstand kommt. Vermutlich versucht tf, das ganze D (und M?) Im Speicher zu speichern.

Wenn ich dies in C/C++ schreiben würde, würde ich höchstwahrscheinlich die Matrixzeilen überlappen, jede Zeile summieren, während ich gehe, und nie die ganze Matrix speichern. Dito die GPU - ich würde die (virtuelle) Matrix unterteilen und die Reduktion in Chunks durchführen.

Gibt es einen Trick, um tf dazu zu bringen, einem mehr stückweisen Verhalten zu folgen und dabei Speicher zu sparen?

Cheers,

Chris

EDIT:

Ein alternativer Ansatz, der mit dem Speicherproblem verkraftet verwenden tf.map_fn:

rowsums = tf.map_fn(lambda i: tf.reduce_sum(tf.sqrt(tf.reduce_sum(tf.pow(i - x,2),1))) , x) 
res = tf.reduce_sum(rowsums) 

So werden nur die rowsums als eine gespeicherte Tensor und nicht die Volldistanzmatrix. Obwohl dieser Ansatz auf der CPU gut funktioniert, kommt es auf der GPU zum Stillstand.

Answer 1

Was hier wirklich benötigt wird (aber noch nicht implementiert ist) ist cwise fusion. Was gerade passiert, ist, dass 2*sqrt(a+b) neuen Tensor für , dann neuen Tensor für sqrt und dann noch einen für 2*sqrt zuweisen wird. PS, können Sie graben, wo der Speicher Speicherzuweisung durch die Untersuchung messages gehen wird (verbose logging brauchen)

Sie Dinge mehr Speicher effizient unter Verwendung von Variablen und assign_add inkrementell Dinge zu aktualisieren, ohne die Schaffung viele Zwischen Tensoren machen könnte. Es gibt eine alternative Formel für die Berechnung "alle paarweise Abstände" here, die einfacher zu diesem Formular zu konvertieren sein können