Emulating nested for loops with scan is slow - numpy

Question

I am trying to simulate nesting for loops by using the scan function, but this is slow. Is there a better way to simulate nesting for loops with Tensorflow? I am not doing this computation with solely numpy so that I can have automatic differentiation.
Specifically, I am convolving over an image with a bilateral filter all while using Tensorflow control ops. To accomplish this, I nested scan() functions, but this leaves me with remarkably poor performance---filtering a small image takes more than 5 minutes.
Is there a better way than nesting scan functions, and how badly am I using Tensorflow control flow operations? I'm interested in general answers more than one specific for my code.
Here is the original, faster code if you want to see it:
def bilateralFilter(image, sigma_space=1, sigma_range=None, win_size=None):
if sigma_range is None:
sigma_range = sigma_space
if win_size is None: win_size = max(5, 2 * int(np.ceil(3*sigma_space)) + 1)
win_ext = (win_size - 1) / 2
height = image.shape[0]
width = image.shape[1]
# pre-calculate spatial_gaussian
spatial_gaussian = []
for i in range(-win_ext, win_ext+1):
for j in range(-win_ext, win_ext+1):
spatial_gaussian.append(np.exp(-0.5*(i**2+j**2)/sigma_space**2))
padded = np.pad(image, win_ext, mode="edge")
out_image = np.zeros(image.shape)
weight = np.zeros(image.shape)
idx = 0
for row in xrange(-win_ext, 1+win_ext):
for col in xrange(-win_ext, 1+win_ext):
slice = padded[win_ext+row:height+win_ext+row,
win_ext+col:width+win_ext+col]
value = np.exp(-0.5*((image - slice)/sigma_range)**2) \
* spatial_gaussian[idx]
out_image += value*slice
weight += value
idx += 1
out_image /= weight
return out_image
This is the Tensorflow version:
sess = tf.InteractiveSession()
with sess.as_default():
def bilateralFilter(image, sigma_space, sigma_range):
win_size = max(5., 2 * np.ceil(3 * sigma_space) + 1)
win_ext = int((win_size - 1) / 2)
height = tf.shape(image)[0].eval()
width = tf.shape(image)[1].eval()
spatial_gaussian = []
for i in range(-win_ext, win_ext + 1):
for j in range(-win_ext, win_ext + 1):
spatial_gaussian.append(np.exp(-0.5 * (i ** 2 +\
j ** 2) / sigma_space ** 2))
# we use "symmetric" as it best approximates "edge" padding
padded = tf.pad(image, [[win_ext, win_ext], [win_ext, win_ext]],
mode='SYMMETRIC')
out_image = tf.zeros(tf.shape(image))
weight = tf.zeros(tf.shape(image))
spatial_index = tf.constant(0)
row = tf.constant(-win_ext)
col = tf.constant(-win_ext)
def cond(padded, row, col, weight, out_image, spatial_index):
return tf.less(row, win_ext + 1)
def body(padded, row, col, weight, out_image, spatial_index):
sub_image = tf.slice(padded, [win_ext + row, win_ext + col],
[height, width])
value = tf.exp(-0.5 *
(((image - sub_image) / sigma_range) ** 2)) *
spatial_gaussian[spatial_index.eval()]
out_image += value * sub_image
weight += value
spatial_index += 1
row, col = tf.cond(tf.not_equal(tf.mod(col,
tf.constant(2*win_ext + 1)), 0),
lambda: (row + 1, tf.constant(-win_ext)),
lambda: (row, col))
return padded, row, col, weight, out_image, spatial_index
padded, row, col, weight, out_image, spatial_index =
tf.while_loop(cond, body,
[padded, row, col, weight, out_image, spatial_index])
out_image /= weight
return out_image
cat = plt.imread("cat.png") # grayscale
cat = tf.reshape(tf.constant(cat), [276, 276])
cat_blurred = bilateralFilter(cat, 2., 0.25)
cat_blurred = cat_blurred.eval()
plt.figure()
plt.gray()
plt.imshow(cat_blurred)
plt.show()

Answer 1

Here is one problem with your code. cols() has a bunch of python globals and you seemed to expect them to be updated at each loop iteration. Please take a look at the TensorFlow tutorial about graph construction and execution. In short, those python globals and their associated code will only be executed at graph construction time, and they are not even in TensorFlow's execution graph. An operation can only be included in the execution graph if it is a tf operator.
Also, it seems that tf.while_loop is better suited for your code than scan.