使用扫描模拟嵌套for循环很慢

时间:2016-06-18 00:41:51

标签: numpy tensorflow

我试图通过使用扫描功能来模拟循环嵌套,但这很慢。有没有更好的方法来使用Tensorflow模拟循环嵌套?我没有单独使用numpy进行这种计算,因此我可以自动区分。

具体来说,我在使用Tensorflow控制操作时使用双边滤波器对图像进行卷积。为了实现这一点,我嵌套了scan()函数,但这使我的性能非常差 - 过滤小图像需要5分钟以上。

是否有比嵌套扫描功能更好的方法,以及使用Tensorflow控制流程操作有多糟糕? 我对我的代码中不止一个特定答案感兴趣。

如果你想看到它,这是原始的,更快的代码:

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

这是Tensorflow版本:

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()

1 个答案:

答案 0 :(得分:1)

以下是您的代码存在的一个问题。 cols()有一堆python全局变量,你似乎期望它们在每次循环迭代时都会被更新。请查看有关图形构造和执行的TensorFlow教程。简而言之,那些python全局变量及其相关代码将仅在图形构造时执行,并且它们甚至不在TensorFlow的执行图中。如果操作只是tf操作符,则操作只能包含在执行图中。

此外,似乎tf.while_loop比扫描更适合您的代码。