Tensorflowでテンソルを複製せずにバッチ内のペアごとの距離を計算しますか？

Question

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​​）との間の二乗距離を計算します各行ベクトル。出力は、要素 'd_ij = dist（x_i、y_j）'を持つサイズ（m、n）の行列です。

問題は、テンソルを複製することで、大量のメモリを消費する大きなバッチと高いディム機能を持つことです。 私はメモリ使用量を増やすことなくこれを実装する別の方法を探していて、最後の距離のテンソルだけを保存します。元のテンソルの二重ループの種類。

Answer 1

一部の線形代数を使用して行列演算に変換できます。あなたはa[i]があなたの元の行列のi番目の行で、

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

あなたはr[i]は、元のi番目の行のノルムを二乗され

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

にそれを書き換えることができDをマトリックス必要なものがありますマトリックス。

あなたは

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) 

としてこれを書くことができ、列ベクトルとして rを扱い、TensorFlowで D

として

D = r - 2 A A' + r' 

を書くことができ、標準broadcasting rulesをサポートするシステムで

Answer 2

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 

ことの一つは、@YaroslavBulatovによってアプローチがそうではないtf.squared_differenceを使用して、このソリューションは、非常に大きなベクトル用のメモリ（OOM）のうち、私を与えることです。だから、私は操作を分解することでメモリフットプリントが小さくなると思っています（私はsquared_differenceがフードの中でより良く扱えると思っていました）。