tensorflow: memory allocation for a 'for' cycle

Question

I am trying to use TensorFlow for calculating minimum Euclidean distance between each column in the matrix and all other columns (excluding itself):

with graph.as_default():
  ...
  def get_diversity(matrix):
       num_rows = matrix.get_shape()[0].value
       num_cols = matrix.get_shape()[1].value
       identity = tf.ones([1, num_cols], dtype=tf.float32)
       diversity = 0

       for i in range(num_cols):
           col = tf.reshape(matrix[:, i], [num_rows, 1])
           col_extended_to_matrix = tf.matmul(neuron_matrix, identity)
           difference_matrix = (col_extended_to_matrix - matrix) ** 2
           sum_vector = tf.reduce_sum(difference_matrix, 0)
           mask = tf.greater(sum_vector, 0)
           non_zero_vector = tf.select(mask, sum_vector, tf.ones([num_cols], dtype=tf.float32) * 9e99)
           min_diversity = tf.reduce_min(non_zero_vector)
           diversity += min_diversity

       return diversity / num_cols
  ...

  diversity = get_diversity(matrix1)

  ...

When I call get_diversity() once per 1000 iterations (on the scale of 300k) it works just fine. But when I try to call it at every iteration the interpreter returns:

W tensorflow/core/common_runtime/bfc_allocator.cc:271] Ran out of memory trying to allocate 2.99MiB.  See logs for memory state.

I was thinking that was because TF creates a new set of variables each time get_diversity() is called. I tried this:

def get_diversity(matrix, scope):
    scope.reuse_variables()
...
with tf.variable_scope("diversity") as scope:
    diversity = get_diversity(matrix1, scope)

But it did not fix the problem.

How can I fix this allocation issue and use get_diversity() with large number of iterations?

I think you should only call get_diversity once. Then every time you want to do the computation you can evaluate the diversity tensor in a session. — Aaron, May 10 '16 at 16:32

score 0 · Answer 1 · edited Sep 22 '17 at 17:08

Assuming you call get_diversity() multiple times in your training loop, Aaron's comment is a good one: instead you can do something like the following:

diversity_input = tf.placeholder(tf.float32, [None, None], name="diversity_input")
diversity = get_diversity(matrix)

# ...

with tf.Session() as sess:
  for _ in range(NUM_ITERATIONS):
    # ...

    diversity_val = sess.run(diversity, feed_dict={diversity_input: ...})

This will avoid creating new operations each time round the loop, which should prevent the memory leak. This answer has more details.

tensorflow: memory allocation for a 'for' cycle

1 Answers1