如何优化这个张量流加权平均定制运算

Question

我正试图实现一个基于本文https://thijsvogels.nl/kpcn/bako2017kpcn.pdf的Tensorflow操作来执行加权平均

op是用权重乘以相邻像素的值来计算图像中像素的平均值

我想寻求任何建议，以优化这段代码，因为目前的实施是相当缓慢

inputs.shape（）是[1，740，1300，3]

weights.shape（）是[1，720，1280，441]

def weighted_average(inputs, weights):
    with tf.name_scope("weighted_average", "weighted_average", [inputs, weights]) as scope:
        in_shape = inputs.get_shape().as_list()
        w_shape = weights.get_shape().as_list()

        n_channels = in_shape[3]
        xs = tf.split(inputs, n_channels, axis=3)

        pad = (in_shape[1] - w_shape[1]) // 2

        kernel_size = pad * 2 + 1

        for index in range(n_channels):
            x = xs[index]

            x_stack = []
            for i in range(kernel_size):
                for j in range(kernel_size):
                    x_stack.append( x[:, i:x.shape[1] - 2 * pad + i, j:x.shape[2] - 2 * pad + j, :] )

            x_stack = tf.concat(x_stack, axis=3)
            x = tf.reduce_sum(tf.multiply(x_stack, weights), axis=3, keep_dims=True)

            xs[index] = x

        return tf.concat(xs, axis=3)

Answer 1

将tf.device('/cpu:0')强制在CPU中计算op，并使用Eigen lib使其速度更快

我想如果用GPU计算的话，可能和所有的张量变换有关

def weighted_averagex(inputs, weights):
    with tf.name_scope("weighted_average", "weighted_average", [inputs, weights]) as scope:
      with tf.device('/cpu:0'):
        in_shape = inputs.get_shape().as_list()
        w_shape = weights.get_shape().as_list()

        n_channels = in_shape[3]
        xs = tf.split(inputs, n_channels, axis=3)

        pad = (in_shape[1] - w_shape[1]) // 2

        kernel_size = pad * 2 + 1

        for index in range(n_channels):
            x = xs[index]

            x_stack = []
            for i in range(kernel_size):
                for j in range(kernel_size):
                    x_stack.append( x[:, i:x.shape[1] - 2 * pad + i, j:x.shape[2] - 2 * pad + j, :] )

            x_stack = tf.concat(x_stack, axis=3)
            x = tf.reduce_sum(tf.multiply(x_stack, weights), axis=3, keep_dims=True)

            xs[index] = x

      return tf.concat(xs, axis=3)