결론부터 말씀드리면, 거리 행렬을 구할 때 X[:, None, :] - X[None, :, :] 식으로 한 번에 브로드캐스팅하면 (N, N, D) 크기의 중간 배열이 그대로 메모리에 올라가서 터집니다. 해결책은 ① 제곱합 전개(내적 트릭)로 중간 배열 자체를 없애거나, ② 행 단위로 쪼개서 필요한 블록만 메모리에 올리는 것입니다. 데이터가 커질수록 numpy 브로드캐스팅에서 메모리가 터지는 지점은 거의 항상 이 (N, N, D) 3차원 중간 배열이고, 아래에서 실제 바이트 수까지 계산해 가며 짚어보겠습니다.
이 글은 numpy의 공식 브로드캐스팅 규칙(버전에 상관없이 numpy 1.x 전반에서 동일하게 적용됨)을 전제로 하고, 거리 행렬 계산에 흔히 쓰는 유클리드 거리를 예시로 듭니다.
왜 브로드캐스팅 거리 행렬에서 메모리가 터질까요
벡터 N개가 있을 때 모든 쌍의 거리를 구하려고 흔히 이렇게 씁니다.
import numpy as np
N, D = 20000, 128
X = np.random.rand(N, D).astype(np.float64)
diff = X[:, None, :] - X[None, :, :] # (20000, 20000, 128)
dist = np.sqrt((diff ** 2).sum(axis=2))
여기서 X[:, None, :]는 (N, 1, D), X[None, :, :]는 (1, N, D)로 브로드캐스팅 규칙에 따라 (N, N, D) 모양으로 확장되면서 뺄셈이 일어납니다. N=20000, D=128일 때 이 중간 배열 하나의 크기를 계산해 보면 20000 × 20000 × 128 × 8바이트(float64) = 약 409.6GB, GiB 단위로는 381GiB에 달합니다.
실제로 이 코드를 돌리면 numpy는 다음과 비슷한 형태의 에러를 냅니다.
MemoryError: Unable to allocate 381. GiB for an array
with shape (20000, 20000, 128) and data type float64N을 2배로 늘리면 필요한 메모리는 4배로 뛰기 때문에, 데이터가 조금만 커져도 금방 수백 GB 단위로 치솟습니다. 문제는 연산 자체가 아니라 중간 단계에서 D축까지 그대로 들고 있는 이 (N, N, D) 배열입니다.
거리 공식을 바꾸면 메모리 사용량이 확 줄어듭니다
가장 근본적인 해법은 공식 자체를 바꾸는 것입니다. 두 벡터 a, b의 유클리드 거리 제곱은 ‖a-b‖² = ‖a‖² + ‖b‖² – 2(a·b)로 전개할 수 있습니다. 이 식을 쓰면 D축을 직접 들고 있을 필요 없이 행렬곱 한 번으로 끝낼 수 있습니다.
sq = np.sum(X ** 2, axis=1) # (N,)
dist_sq = sq[:, None] + sq[None, :] - 2 * X @ X.T
dist_sq = np.maximum(dist_sq, 0) # 부동소수점 오차로 생기는 음수 방지
dist = np.sqrt(dist_sq)
여기서 가장 큰 중간 배열은 (N, N) 하나뿐입니다. N=20000이면 20000 × 20000 × 8바이트 = 약 2.98GiB로, 앞서 계산한 381GiB와는 비교가 안 되게 줄어듭니다. 게다가 X @ X.T는 BLAS로 처리되는 행렬곱이라 단순 반복문보다 빠릅니다.
다만 공짜는 아닙니다. 두 벡터가 거의 같을 때 ‖a‖² + ‖b‖²와 2(a·b)가 비슷한 크기라서 빼기 과정에서 부동소수점 오차로 아주 작은 음수가 나올 수 있습니다. 이 값을 그대로 np.sqrt에 넣으면 NaN이 나오기 때문에 np.maximum(dist_sq, 0)으로 음수를 걷어내는 과정이 꼭 필요합니다.
행 단위 청크 처리로 피크 메모리 잡기
공식을 바꾸기 애매한 거리 지표(예: 맨해튼 거리처럼 내적 전개가 안 되는 경우)를 써야 한다면, 전체를 한 번에 계산하지 않고 행을 나눠서 처리하는 방식이 현실적입니다.
chunk = 500
X32 = X.astype(np.float32)
dist = np.empty((N, N), dtype=np.float32)
for i in range(0, N, chunk):
block = X32[i:i + chunk] # (chunk, D)
d = block[:, None, :] - X32[None, :, :] # (chunk, N, D)만 메모리에 올림
dist[i:i + chunk] = np.sqrt((d ** 2).sum(axis=2))
이렇게 하면 한 번에 메모리에 올라가는 중간 배열은 (chunk, N, D) 크기뿐입니다. chunk=500, N=20000, D=128, float32 기준으로 계산하면 500 × 20000 × 128 × 4바이트 = 약 4.77GiB이고, 반복마다 이 공간을 재사용하기 때문에 전체 계산이 끝날 때까지 이 수치 이상으로 메모리가 늘어나지 않습니다. float64 대신 float32를 쓴 것도 메모리를 절반으로 줄이는 데 한몫합니다.
chunk 크기는 직접 조절하는 값이라, 가진 메모리에 맞춰 500, 1000, 2000 등으로 바꿔가며 OOM이 안 나는 선을 찾으면 됩니다. 단점은 파이썬 for문이 추가되면서 완전히 벡터화된 한 줄 연산보다는 약간의 오버헤드가 생긴다는 점입니다.
scipy와 sklearn, 상황에 맞게 골라 쓰세요
직접 구현하지 않고 라이브러리에 맡기는 선택지도 있습니다. scipy.spatial.distance.cdist는 내부적으로 C로 구현돼 있어 (N, N, D) 같은 파이썬 레벨 중간 배열을 만들지 않고 결과 (N, N) 행렬만 반환합니다. 데이터셋이 (N, N) 결과 자체도 못 들고 있을 만큼 크다면 sklearn.metrics.pairwise_distances_chunked처럼 결과를 청크 단위 제너레이터로 흘려주는 함수를 쓰는 편이 낫습니다.
| 방법 | 추가로 필요한 메모리 | 속도 특성 | 주의할 점 |
|---|---|---|---|
| 브로드캐스팅 한 번에 계산 | (N, N, D) 전체 | 코드는 가장 짧지만 메모리 요구량이 가장 큼 | N, D가 조금만 커져도 OOM |
| 제곱합 전개(내적 트릭) | (N, N)만 | 행렬곱(BLAS) 덕분에 빠름 | 부동소수점 오차로 음수 클리핑 필요, 유클리드 거리 전용 |
| 행 단위 청크 반복 | (chunk, N, D) | 벡터화 유지, chunk 조절 가능 | 파이썬 루프 오버헤드, 어떤 거리 지표든 적용 가능 |
| scipy.spatial.distance.cdist | (N, N)만 | C 구현이라 빠름 | 지원하는 거리 지표가 정해져 있음 |
| sklearn pairwise_distances_chunked | 청크 단위 | 디스크/스트리밍 처리에 적합 | 결과를 한 번에 안 주고 제너레이터로 받음, 추가 의존성 |
청크로 쪼개도 메모리가 부족할 때가 있을까요
청크 처리는 ‘중간 배열’의 폭발을 막아줄 뿐, 최종 결과물인 (N, N) 거리 행렬 자체의 크기까지 줄여주지는 않습니다. N이 20만 개로 커진다면 float32로 저장해도 결과 행렬만 200000 × 200000 × 4바이트 = 약 149GiB가 되어, 청크 처리로 중간 계산을 아무리 잘 쪼개도 출력 배열 자체를 메모리에 다 들고 있을 수 없습니다.
이 구간에 들어서면 접근 자체를 바꿔야 합니다. 전체 쌍별 거리를 다 구하는 대신 np.memmap으로 결과를 디스크에 두고 블록 단위로 써 내려가거나, 애초에 모든 쌍의 거리가 아니라 각 점의 k-최근접 이웃만 필요하다면 FAISS 같은 근사 최근접 이웃 라이브러리로 문제 자체를 바꾸는 편이 현실적입니다. “메모리를 더 아끼는 코드”로는 못 넘는 지점이 있다는 뜻입니다.
지금 코드의 거리 행렬 한 줄만 바꿔보세요
정리하면, 거리 행렬에서 numpy 브로드캐스팅으로 메모리가 터지는 원인은 거의 항상 D축을 들고 있는 (N, N, D) 중간 배열입니다. 유클리드 거리라면 제곱합 전개로 (N, N)만 남기는 게 가장 간단한 해법이고, 다른 거리 지표를 써야 한다면 행 단위 청크 반복이나 scipy의 cdist로 바꾸는 게 현실적입니다. N이 수십만 단위로 커져서 결과 행렬 자체가 메모리를 넘어선다면, 메모리 절약 기법 대신 memmap이나 근사 최근접 이웃 같은 다른 접근을 검토해야 합니다.
오늘 당장 할 수 있는 건, 지금 쓰고 있는 거리 행렬 코드에서 diff = X[:, None, :] - X[None, :, :] 한 줄을 찾아서 제곱합 전개 버전으로 바꿔보고, N과 D를 넣어 실제로 몇 GiB가 줄어드는지 직접 계산해 보는 것입니다.
