Управление памятью Tensorflow - chunking?

Question

У меня проблема с вычислением сокращения функции сети, представленной большой (200000x200000) матрицей, сгенерированной как матрица расстояний между парами точек.

Минимальный пример, вход X A 200000x2 Numpy массив декартовых координат:

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) 

, работающие на процессоре, памяти (16 Гб на моем MBP) быстро превысил предложение, и система перемалывает к остановке. Предположительно tf пытается хранить в памяти все D (и M?).

Если бы я писал это на C/C++, я бы, скорее всего, переместил строки матрицы, суммируя каждую строку, когда я иду, и никогда не сохраняю всю матрицу. Тоже GPU - я бы разделил (виртуальную) матрицу и выполнил сокращение кусков.

Есть ли трюк, чтобы получить tf, чтобы следить за более частым поведением, экономя память?

Приветствия,

Криса

РЕДАКТИРОВАТЬ:

Альтернативный подход, который справляется с проблемой памяти является использование 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) 

Таким образом, только rowsums сохраняются в виде тензор, а не матрицу полного расстояния. Однако, хотя этот подход хорошо работает на процессоре, он останавливается на GPU.

Answer 1

Что действительно нужно здесь (но еще не реализовано) - это слияние. Что происходит сейчас, так это то, что 2*sqrt(a+b) выделит новый тензор для a+b, затем новый тензор для sqrt, а затем еще один для 2*sqrt. PS, вы можете выкопать, где память собирается путем анализа распределения памяти messages (необходимо verbose logging)

Вы могли бы сделать более эффективной память с помощью переменных и assign_add пошагово обновлять вещи, не создавая множество промежуточных тензоров. Существует альтернативная формула для вычисления «всех попарных расстояний» here, которые могут быть легко преобразованы в эту форму.