Tensorflow 메모리 관리 - 청킹?

Question

포인트 쌍 사이의 거리 행렬로 생성 된 큰 (200000x200000) 행렬로 표시되는 네트워크의 기능 축소를 계산하는 데 문제가 있습니다. 신속 초과 가입하고 시스템 중단 갈기 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기가바이트)

최소한의 예는, 입력 직교 좌표 200000x2 NumPy와 배열하기 x. 아마도 tf는 D (및 M?) 전체를 메모리에 저장하려고 시도하고있을 것입니다.

C/C++로이 글을 쓰고 있다면 행렬 전체를 루프 처리하고 전체 행렬을 저장하지 않으므로 각 행을 합산하는 것이 좋습니다. GPU와 동등한 - (가상) 행렬을 세분화하고 청크 감소를 수행합니다.

tf가 메모리를 절약하면서 더 많은 청크와 같은 동작을 수행하는 속임수가 있습니까?

건배

크리스

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에 대해 새 Tensor를 할당 한 다음 sqrt에 대해 새 텐서를 할당 한 다음 2*sqrt으로 또 다른 텐서를 할당합니다. 메모리가 메모리 할당 messages을 검사하여 어디로 PS, 당신은

당신은 점차적으로 많은 중간 텐서를 만들지 않고 일을 업데이트 할 변수와 assign_add를 사용하여 일 더 많은 메모리를 효율적으로 만들 수 (verbose logging 필요) 팔 수 있습니다. "모든 쌍 거리 (pairwise distances)"를 계산하기위한 다른 공식이 있습니다.이 형식으로 변환하기가 더 쉬울 수도 있습니다.