Tensorflowメモリ管理 - チャンク？

Question

私は、ポイントのペア間の距離マトリックスとして生成された大きな（200000x200000）マトリックスによって表されるネットワークの機能の低下を計算することに問題があります。

最小限例えば、デカルト座標の200000x2 numpyの配列Xの入力：CPU上で実行

x = tf.constant(X[:,0], shape=[X.shape[0],1]) 
y = tf.constant(X[:,1], shape=[X.shape[0],1]) 
dx = x - tf.transpose(x) 
dy = y - tf.transpose(y) 
D = tf.sqrt(dx*dx + dy*dy) 
M = 0.1 * 5.0/tf.pow(4.0 + D, 1.5) 
res = tf.reduce_sum(betaM) 

は、メモリ（私のMBPに16ギガバイト）が迅速にオーバーサブスクライブされ、システムが停止する磨きます。多分、tfはD（とM？）の全体をメモリに格納しようとしています。

これをC/C++で書いていたのであれば、行列の行全体をループし、各行を合計して行列全体を決して格納しないようにしています。 GPUと同上 - 私は（仮想）マトリックスを細分し、チャンクの縮小を実行します。

メモリを節約しながら、より多くのチャンクに従うようにするには、トリックがありますか？

乾杯、

クリス

EDIT：

メモリの問題に対処する別のアプローチは、tf.map_fnを使用することである。

rowsums = tf.map_fn(lambda i: tf.reduce_sum(tf.sqrt(tf.reduce_sum(tf.pow(i - x,2),1))) , x) 
res = tf.reduce_sum(rowsums) 

したがってのみrowsumsは次のように格納されていますテンソルであり、全距離行列ではない。しかし、このアプローチはCPU上でうまく機能しますが、GPU上で停止するようになります。

Answer 1

ここで本当に必要なのは（まだ実装されていない）cwise fusionです。今のところ、2*sqrt(a+b)はa+bに新しいテンソルを割り当て、sqrtに新しいテンソル、次に2*sqrtに新しいテンソルを割り当てます。 PS、あなたはメモリがメモリ割り当てmessages（verbose loggingを必要とする）

を調べることによって起こっている場所をインクリメンタルに多くの中間テンソルを作成せずに物事を更新するために、変数とassign_addを使用して、物事より多くのメモリを効率的に作ることができます掘ることができます。 「すべてのペアごとの距離」を計算する代わりの式がありますhereこのフォームに変換する方が簡単かもしれません。