每个张量的Tensorflow梯度

Question

I have a network that takes as input an Nx3 matrix and produces an N-dimensional vector. Let's say batch size is 1 and N=1024, so the output would have the shape (1,1024) . I want to compute the gradients for every dimension of the output, with respect to the input. That is, dy/dx for every y . However tensorflow's tf.gradients computes d sum(y)/dx , aggregate. I know there's no straightforward way to compute the gradients for every output dimension, so I finally decided to run tf.gradients 1024 times, because I only have to do this once in the project, and never again.

So I do this:

start = datetime.datetime.now()
output_code_split = tf.split(output_code,1024)
#output shape = (1024,)
grad_ops = []
for i in range(1024):
    gr = tf.gradients(output_code_split[i],input)
    #output shape = (1024,1,16,1024,3) , where 16= batch size

    gr = tf.reduce_mean(gr,[0,1,2,3])
    #output shape = (1024,)

    grad_ops.append(gr)
    present = datetime.datetime.now()
    print(i,(present-start).seconds,flush=True)
    #prints time taken to finish previous computation.
    start = datetime.datetime.now()

When the code started running, the time between two iterations was 4 seconds, so I figured it'll run for roughly 4096 seconds. However, as the number of iterations increase, the time taken for subsequent runs keeps increasing. The gap, which was 4 seconds when the code started, eventually got to 30 seconds after about 500 iterations, which is too much.

Is the list holding the gradient ops grad_ops growing bigger and occupying more memory. I'm unfortunately not in the position to do a detailed memory profiling of this code..Any ideas about what causes the iteration time to blow up as time goes on?

(Note that in the code, I'm only creating the gradient ops and not actually evaluating them. That part comes later, but my code doesn't reach there on account of the extreme slowdown mentioned above)

Thanks.

Answer 1

What blows up your execution time is that you define a new operation on the graph in every iteration of your for loop. Every call to tf.gradient and tf.reduce_mean pushes a new node onto the graph. Then it needs to recompile to be run. What should actually work for you is to use tf.gather with an int32 placeholder, which supplies the dimension to your gradient operation. So something like this:

idx_placeholder = tf.placeholder(tf.int32, shape=(None,))
grad_operation = tf.gradients(tf.gather(output_code_split, idx_placeholder))
for i in range(1024):
      sess.run(grad_operation, {idx_placeholder: np.array([i])})