Evaluating a tensor partially

Question

Suppose we have a convolution operation such as this:

y = tf.nn.conv2d( ... )

Tensorflow allows you to evaluate a part of a tensor, e.g.:

print(sess.run(y[0]))

When we evaluate a tensor partially like above, which one of the following is correct?

1. TF runs the whole operation, i.e. it computes y completely, and then returns the value of y[0]
2. TF runs only the necessary operations to compute y[0].

Answer 1

I set up a small sample program:

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # forcing to run on the CPU
import tensorflow as tf

def full_array(sess, arr):
  sess.run(arr)[0]
  
def partial_array(sess, arr):
  sess.run(arr[0])
  
sess = tf.Session()
arr = tf.random_uniform([100])
arr = arr + tf.random_uniform([100])

These are my results:

%timeit partial_array(sess, arr)
100 loops, best of 3: 15.8 ms per loop

%timeit full_array(sess, arr)
10000 loops, best of 3: 85.9 µs per loop

It seems from the timings that the partial run is actually much slower than the full run (which is confusing to me to be honest...)

With these timings, I'd exclude alternative 1), since I would expect the timing to be approximately the same in the two functions if that were true.

Given my simplified test code, I'd lean towards the idea that the logic to figure out what part of the graph needs to run to satisfy the tensor slice is the cause of the performance difference, but I don't currently have a proof for that.

Update:

I also ran a similar test with a convolution op instead of an add (which I thought might be excessively simple an example):

def full_array(sess, arr):
  return sess.run(arr)[0]
  
def partial_array(sess, arr):
  return sess.run(arr[0])
  
sess = tf.Session()
arr = tf.random_uniform([1,100,100,3])
conv = tf.nn.conv2d(arr, tf.constant(1/9, shape=[3,3,3,6]), [1,1,1,1], 'SAME')

The results, however, are consistent with the previous ones:

%timeit full_array(sess, conv)
1000 loops, best of 3: 949 µs per loop

%timeit partial_array(sess, conv)
100 loops, best of 3: 20 ms per loop