首页 文章

在Tensorflow中计算批量中的成对距离而不复制张量?

提问于
浏览
22

我想计算Tensorflow中一批特征的成对平方距离 . 我有一个简单的实现使用和*操作平铺原始张量:

def pairwise_l2_norm2(x, y, scope=None):
    with tf.op_scope([x, y], scope, 'pairwise_l2_norm2'):
        size_x = tf.shape(x)[0]
        size_y = tf.shape(y)[0]
        xx = tf.expand_dims(x, -1)
        xx = tf.tile(xx, tf.pack([1, 1, size_y]))

        yy = tf.expand_dims(y, -1)
        yy = tf.tile(yy, tf.pack([1, 1, size_x]))
        yy = tf.transpose(yy, perm=[2, 1, 0])

        diff = tf.sub(xx, yy)
        square_diff = tf.square(diff)

        square_dist = tf.reduce_sum(square_diff, 1)

        return square_dist

该函数将两个大小为(m,d)和(n,d)的矩阵作为输入,并计算每个行向量之间的平方距离 . 输出是大小为(m,n)的矩阵,其元素为'd_ij = dist(x_i,y_j)' .

问题是我有一个大批量和高昏暗的功能'm,n,d'复制张量消耗了大量的内存 . 我正在寻找另一种方法来实现它,而不增加内存使用量,只是存储最终的距离张量 . 一种双循环原始张量 .

4 回答

  • 44

    您可以使用一些线性代数将其转换为矩阵运算 . 注意你需要的是矩阵 D ,其中 a[i] 是原始矩阵的 i

    D[i,j] = (a[i]-a[j])(a[i]-a[j])'
    

    你可以把它重写成

    D[i,j] = r[i] - 2 a[i]a[j]' + r[j]
    

    其中 r[i] 是原始矩阵的 i 行的平方范数 .

    在支持标准broadcasting rules的系统中,您可以将 r 视为列向量并将 D 写为

    D = r - 2 A A' + r'
    

    在TensorFlow中,您可以将其写为

    A = tf.constant([[1, 1], [2, 2], [3, 3]])
    r = tf.reduce_sum(A*A, 1)
    
    # turn r into column vector
    r = tf.reshape(r, [-1, 1])
    D = r - 2*tf.matmul(A, tf.transpose(A)) + tf.transpose(r)
    sess = tf.Session()
    sess.run(D)
    

    结果

    array([[0, 2, 8],
           [2, 0, 2],
           [8, 2, 0]], dtype=int32)
    
  • 3

    使用 squared_difference

    def squared_dist(A): 
        expanded_a = tf.expand_dims(A, 1)
        expanded_b = tf.expand_dims(A, 0)
        distances = tf.reduce_sum(tf.squared_difference(expanded_a, expanded_b), 2)
        return distances
    

    我注意到的一件事是,使用 tf.squared_difference 的这个解决方案让我为非常大的向量提供内存(OOM),而@YaroslavBulatov的方法却没有 . 因此,我认为分解操作会产生更小的内存占用(我认为 squared_difference 会更好地处理它) .

  • 0

    这是坐标 AB 的两个张量的更通用的解决方案:

    def squared_dist(A, B):
      assert A.shape.as_list() == B.shape.as_list()
    
      row_norms_A = tf.reduce_sum(tf.square(A), axis=1)
      row_norms_A = tf.reshape(row_norms_A, [-1, 1])  # Column vector.
    
      row_norms_B = tf.reduce_sum(tf.square(B), axis=1)
      row_norms_B = tf.reshape(row_norms_B, [1, -1])  # Row vector.
    
      return row_norms_A - 2 * tf.matmul(A, tf.transpose(B)) + row_norms_B
    

    请注意,这是方形距离 . 如果要将其更改为欧几里德距离,请对结果执行 tf.sqrt . 如果你想这样做,不要忘记添加一个小常量来补偿浮点不稳定性: dist = tf.sqrt(squared_dist(A, B) + 1e-6) .

  • 9

    如果您想要计算其他方法,那么更改tf模块的顺序 .

    def compute_euclidean_distance(x, y):
        size_x = x.shape.dims[0]
        size_y = y.shape.dims[0]
        for i in range(size_x):
            tile_one = tf.reshape(tf.tile(x[i], [size_y]), [size_y, -1])
            eu_one = tf.expand_dims(tf.sqrt(tf.reduce_sum(tf.pow(tf.subtract(tile_one, y), 2), axis=1)), axis=0)
            if i == 0:
                d = eu_one
            else:
                d = tf.concat([d, eu_one], axis=0)
    return d
    

相关问题