모델링 코드를 작성하거나 전처리를 진행할 때 가장 중요한 것은 코드의 "효율성"이다.
이 때 가장 많이 활용하는 것이 바로 vectorization(벡터화)인데, 아직도 나를 포함하여 많은 학생 및 연구자들은 매우 비효율적인 무한 for 루프 코딩 스타일에서 벗어나지 못하고 있는 것이 현실이다...
따라서 이번 글을 통해 개념을 확실히 다져보도록 하자!
(혹시 오개념이 있다면 댓글로 편하게 지적해주세요)
1. 벡터화가 뭐 그렇게 중요한가..?
> 네!!! 중요해요!!!!!
많이들 벡터화=시간복잡도 감소 라고 생각하지만, 대부분의 경우 알고리즘적으로 Big-O는 동일하다.
예를 들어 N개 원소에 y = x*2 + 1을 적용하면 for문과 벡터화 모두 O(N)이다.
그럼에도 불구하고 벡터화가 빠른 이유는,
- for문은 파이썬이 N번 돌아간다. (해석기 overhead + 동적 타입 + 함수 호출 및 bound check etc.)
- 그러나 벡터화는 C/Fortran(NumPy) 또는 C++/CUDA(PyTorch)로 작성된 내부 루프가 한 번에 돌아간다.
즉, 파이썬은 1번만 호출하고 내부에서 빠른 루프가 돌아가는 구조로서 같은 Big-O를 훨씬 작은 상수로 실행하는 기술이다!
내부 루프가 이렇게 빠른 이유는 다음과 같이 크게 3가지로 정리할 수 있다.
1. 연속 메모리 & 캐시 효율
- ndarray / tensor는 보통 데이터가 연속된 메모리 블록에 들어있다. CPU는 연속 메모리를 읽을 때 캐시가 잘먹어서 훨씬 속도가 빠르다.
- 반면에 파이썬의 list는 원소가 연속된 숫자 덩어리가 아니라, "파이썬 객체"에 대한 "pointer들의 배열"이라 메모리 접근이 산발적이고 느리다.
2. SIMD/멀티스레드와 같은 low-level 최적화
- NumPy는 내부적으로 BLAS 같은 라이브러리나 벡터 명령(SSE/AVX)을 쓰는 경우가 많고, PyTorch도 CPU에서는 벡터화/스레딩, GPU에서는 CUDA 커널로 대량 병렬 처리를 진행한다.
3. GPU에서는 "kernel launch"가 핵심
- GPU 텐서 연산은 보통 파이썬 for문으로 작은 연산을 여러 번 호출하면, 커널을 계속 띄워서 매우 비효율적이다.
- 반면에 큰 tensor 연산을 한 번 호출하면, GPU가 대량 병렬로 한 번에 처리해준다.
따라서 특히 PyTorch로 코딩할 때 벡터화는 사실상 필수라고 볼 수 있는 것!!
2. 벡터화 시작해보기
벡터화를 너무 어렵게 생각하지 않아도 된다.
이름만 거창할 뿐, 사실상 shape 조작 + 축 연산이기 때문에, '원소'가 아닌 '축(axis)'과 'shape'로 생각하면 쉽다.
벡터화를 내 코드에 적용하기 위해서는 다음 3가지 질문들을 반복적으로 짚어나가면 좋다.
- 내 데이터 shape이 무엇인가? (ex. (B, T, F)?, (N, D)?)
- 어느 축을 기준으로 연산해야 하는가? (sum/mean/max가 어느 dim인지)
- shape이 안 맞는 경우 어떻게 맞춰야 할까? (broadcast를 위한 unsqueeze/reshape)
간단한 예시 코드로 이해해보자.
ex) 배치(B) 별로 feature(D) 평균 빼기
- X : (B, D)
- mu: (D,) 또는 (1, D)
- X - mu 는 broadcast로 자동 확장
mu = X.mean(axis=0) # (D,)
X_centered = X - mu # (B,D) - (D,) -> (B,D)
참고로 앞으로 벡터화 코드에서 축(axis)를 많이 다루게 될 텐데, axis=0은 열 / axis=1은 행임을 확실하게 숙지해두고 넘어가면 헷갈리지 않을 것이다!
위 코드를 예로 들자면 만약 X가 2행 3열(2,3)의 배열이라고 할 때, X.mean(axis=0)라고 하면 각 열별로 평균을 낸다는 의미이므로 총 3개의 열이 각각 평균을 내어, 최종적으로 (3,)의 shape을 가진 mu가 만들어진다.
X가 (2, 3)이고, mu가 (3,)이므로 이 둘을 뺄 때 자동으로 broadcast가 되어 X의 첫 번째 행에서도 mu가 차감되고, 두 번째 행에서도 mu가 차감되어 결과적으로 (2, 3) 형태를 유지한 X_centered가 만들어지는 것이다.
*broadcasting은 행과 열 중 하나 이상이 같아야 이루어진다! 혹시 아직 broadcasting을 모른다면 그것부터 공부하고 다시 읽는 것을 추천합니다... 딱히 어려운 개념은 아닌데, 쉽게 말해 서로 다른 차원(shape)을 가진 두 tensor 간의 연산에서 작은 tensor의 shape을 "확장"해서 상대 tensor와 연산이 가능하도록 맞춰주는 작업이라고 보면 된다~
3. 벡터화 예제 코드
벡터화를 제대로 연습하기 위해서는 다음 단계대로 접근해보면 좋다.
(1) Broadcasting/Indexing(masking)
(2) 집계 (sum/mean)
(3) gather/scatter/unfold 같은 tensor 조작
사실 함수별로 더 많은 내용이 있지만.. 분량상 다 작성하지는 못했고, 핵심이 되는 메소드 위주로 작성했기에 더 알아보고 싶다면 구글링을 통해 더 공부해보시길!
3-1. NumPy - 배열 벡터화
numpy의 벡터화는 기본적으로 np.ndarray를 대상으로 한다.
파이썬 리스트에서는 원소별 연산이 아니라 "리스트 복제"가 발생함
import numpy as np
x_list = [1, 2, 3]
print(x_list*2) # 파이썬 리스트 - 리스트 복제
# [1, 2, 3, 1, 2, 3]
x = np.array(x_list)
print(x*2) # numpy 배열 - 원소별 연산
# [2 4 6]
1) ufunc: 원소별 연산을 for문 없이 구현
- numpy의 ufunc(universal function)는 배열 원소별 연산을 C level에서 처리하게 해주는 기본 도구이므로 for문보다 빠르다.
# 1) 수학/비교/논리 연산은 대부분 ufunc
x = np.array([1.0, 4.0, 9.0])
print(np.sqrt(x)) # [1. 2. 3.]
print(np.log(x)) # [0. 1.38629436 2.19722458]
print(x > 3) # [False True True]
print(np.where(x > 3, 1, 0)) # [0 1 1]# ReLU 함수 구현
x = np.array([-3, -1, 0, 2, 5])
y = np.maximum(x, 0) # if 없이 벡터화
print(y) # [0 0 0 2 5]
2) Broadcasting: 모양(shape)이 달라도 자동으로 맞춰 연산
- 핵심 규칙: 뒤 차원부터 비교해서
- 두 차원이 같거나
- 둘 중 하나가 1이면
호환됨!
# (N, D) ± (D,)
X = np.random.randn(5, 3) # (5, 3)
print(X)
mu = X.mean(axis=0) # (3,)
print(mu)
X_centered = X - mu # mu가 (5, 3)으로 broadcast
print(X_centered)
print(X.shape, mu.shape, X_centered.shape)
'''
[[-0.91009788 -0.07388561 0.37987671]
[-0.4815046 0.29419912 0.2965512 ]
[-1.15978472 -1.01129191 -2.12807028]
[-0.06942 -0.19358232 -1.09222469]
[-1.61804038 1.7773785 -0.79827284]]
[-0.84776951 0.15856356 -0.66842798]
[[-0.06232836 -0.23244917 1.04830469]
[ 0.36626491 0.13563556 0.96497918]
[-0.3120152 -1.16985547 -1.4596423 ]
[ 0.77834952 -0.35214587 -0.42379671]
[-0.77027087 1.61881494 -0.12984486]]
(5, 3) (3,) (5, 3)
'''# (N,)을 (N,1)로 만들어서 (N, M)과 연산하기
a = np.array([1, 2, 3]) # (3,)
b = np.array([10, 20, 30, 40]) # (4,)
# a[: None] => (3, 1), b[None, :] => (1, 4)
grid = a[:, None] + b[None, :]
print(grid.shape)
print(grid)
'''
(3, 4)
[[11 21 31 41]
[12 22 32 42]
[13 23 33 43]]
'''
3. Indexing: 슬라이싱 + boolean mask + fancy index
# 1. 기본 슬라이싱
X = np.arange(20).reshape(4, 5) # (4, 5)
print(X)
print(X[:, 2:4])
'''
[[ 0 1 2 3 4]
[ 5 6 7 8 9]
[10 11 12 13 14]
[15 16 17 18 19]]
[[ 2 3]
[ 7 8]
[12 13]
[17 18]]
'''
# 2. boolean mask로 필터링/치환
X = np.array([[-1, 2, 3],
[4, -5, 6]])
mask = (X < 0)
print(mask)
X2 = X.copy()
X2[mask] = 0
print(X2)
'''
[[ True False False]
[False True False]]
[[0 2 3]
[4 0 6]]
'''
# 3. 조건 + 조건 (and, or) 조합
x = np.array(np.arange(1, 10))
mask = (x >= 3) & (x <= 7) # 파이썬의 and가 아니라 &를 사용!!
print(x[mask]) # [3 4 5 6 7]
# 4. Fancy indexing: 원하는 인덱스만 뽑기
x = np.array([10, 20, 30, 40, 50])
idx = np.array([0, 3, 3, 1])
print(x[idx]) # [10 40 40 20]
4) Axis(축) 기반 집계: sum/mean/max로 for문 제거
- for문 없이 n차원 배열에서 함수 적용/axis 이해하기
- axis=0: "행 방향으로 줄여서(=열별 집계)"
- axis=1: "열 방향으로 줄여서(=행별 집계)"
X = np.arange(12).reshape(3, 4)
print(X.sum(axis=0)) # 열 합: (4,)
print(X.sum(axis=1)) # 행 합: (3,)
print(X.mean(axis=1)) # 행 평균: (3,)
print(X.max(axis=0)) # 열 최대값: (4,)
'''
[12 15 18 21]
[ 6 22 38]
[1.5 5.5 9.5]
[ 8 9 10 11]
'''
5) reshape/transpose: 모양을 바꿔서 벡터화가 되게 만들기
- reshape는 데이터는 바꾸지 않고 모양만 바꿈 (원소 개수는 유지)
- broadcasting이 안 될 때 (N,)을 (N,1)로 reshape 하는 게 전형적인 해결책
x = np.arange(12)
A = x.reshape(3, 4)
B = A.T
print(A.shape, B.shape) # (3, 4) (4, 3)
6) 기타
# 1. 조건 분기 벡터화 - np.where / np.select
x = np.array([-2, -1, 0, 1, 2])
y = np.where(x > 0, x, 0) # if x > 0 return x else return 0
print(y) # [0 0 0 1 2]
# 2. 행마다 다른 인덱스로 뽑기 - np.take_along_axis
# ex) batch별 argmax 위치 값 뽑기
X = np.array([[1, 9, 3],
[7, 2, 8]])
idx = np.array([[1],
[2]]) # (N, 1)
out = np.take_along_axis(X, idx, axis=1) # (N, 1)
print(out.ravel()) # [9 8]
# 3. 구간(bin) 매핑을 for문 없이 - np.searchsorted / np.digitize
edges = np.array([0, 10, 20, 30])
x = np.array([3, 17, 27, 29])
bin_id = np.digitize(x, edges) - 1
print(bin_id) # [0 1 2 2]
# 4. 다차원 곱/축 합을 한 줄로 - np.einsum / matmul / tensordot
# pairwise dot, 가중합, 배치 행렬곱 등에서 for문 제거
A = np.random.randn(5, 3)
B = np.random.randn(5, 3)
print(A)
print(B)
dot = np.einsum('nd,nd->n', A, B) # 각 행의 dot: (5,)
# 각 행(row)을 독립적인 벡터로 보고 대응하는 행끼리 dot product 수행
# 두 입력 배열의 행 수와 동일한 길이의 1차원 배열 반환
print(dot)
'''
[[ 0.21946471 0.91494708 -0.02927509]
[-1.07513621 0.34412369 -0.32883659]
[ 1.07092064 0.71124095 -0.17486011]
[-0.93450915 -0.84330364 -0.73102974]
[-1.08903744 1.49715056 -0.78400969]]
[[ 0.01897095 0.8727142 -0.43942211]
[ 0.46391088 -2.26759528 1.71744067]
[ 0.17665819 -0.12149574 0.3904708 ]
[-0.96478881 1.71411718 -1.49935351]
[ 0.31520855 0.11252911 -1.40646561]]
[ 0.81551489 -1.84385799 0.03449639 0.55215473 0.92788177]
'''
# sum보다 einsum이 더욱 효율적이다!
vec = np.random.rand(100, 100)
print("np.sum:", np.sum(vec))
%timeit np.sum(vec)
print("np.einsum:", np.einsum('ij->', vec))
%timeit np.einsum('ij->', vec)
'''
np.sum: 5009.025489722975
2.14 μs ± 5.44 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
np.einsum: 5009.0254897229715
1.46 μs ± 6.62 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
'''
# 5. 슬라이딩 윈도우 - sliding_window_view
from numpy.lib.stride_tricks import sliding_window_view
x = np.arange(10)
w = sliding_window_view(x, window_shape=4)[::2] # stride 2
print(w.shape)
print(w)
'''
(4, 4)
[[0 1 2 3]
[2 3 4 5]
[4 5 6 7]
[6 7 8 9]]
'''
3-2. PyTorch - 텐서 벡터화
PyTorch는 Tensor = ndarray + autograd + GPU이다!
1. Broadcasting + 기본 원소 연산
- numpy와 완전 동일하며, unsqueeze/None 인덱싱이 자주 등장한다.
import torch
X = torch.randn(32, 128) # (B, D)
mu = X.mean(dim=0) # (D,)
X_centered = X - mu # broadcast
print(X.shape) # torch.Size([32, 128])
print(mu.shape) # torch.Size([128])
print(X_centered.shape) # torch.Size([32, 128])# 원하는 축에 차원 추가하기 - None / unsqueeze
a = torch.arange(3) # (3,)
b = torch.arange(4) # (4,)
grid = a[:, None] + b[None, :] # (3, 4)
print(a)
print(b)
print(grid)
'''
tensor([0, 1, 2])
tensor([0, 1, 2, 3])
tensor([[0, 1, 2, 3],
[1, 2, 3, 4],
[2, 3, 4, 5]])
'''
A0 = torch.unsqueeze(a, dim=0)
print(A0)
A1 = torch.unsqueeze(a, dim=1)
print(A1)
'''
tensor([[0, 1, 2]])
tensor([[0],
[1],
[2]])
'''
2. 마스킹 / 조건 분기 : where, masked_fill, boolean indexing
- 전처리 및 loss 계산에 매우 많이 사용됨!
x = torch.tensor([-2., -1., 0., 1., 2.])
relu = torch.clamp(x, min=0) # tensor([0., 0., 0., 1., 2.])
relu2 = torch.where(x > 0, x, torch.zeros_like(x)) # tensor([0., 0., 0., 1., 2.])X = torch.randn(4, 3)
mask = X < 0
X2 = X.masked_fill(mask, 0.0) # 음수만 0
print(X)
print(mask)
print(X2)
'''
tensor([[ 0.5927, 2.1155, 0.6215],
[-0.2331, 1.5869, 0.6540],
[-1.1697, 1.5475, 1.3878],
[ 0.5975, 0.2010, -0.4702]])
tensor([[False, False, False],
[ True, False, False],
[ True, False, False],
[False, False, True]])
tensor([[0.5927, 2.1155, 0.6215],
[0.0000, 1.5869, 0.6540],
[0.0000, 1.5475, 1.3878],
[0.5975, 0.2010, 0.0000]])
'''
3. 축 집계(reduction): sum/mean/max/... + keepdim
- batch loss, 통계량, score 계산 등에 사용된다.
X = torch.randn(8, 5)
row_sum = X.sum(dim=1) # (8,)
col_mean = X.mean(dim=0) # (5,)
mx = X.max(dim=1).values # (8,)
print(row_sum.shape) # torch.Size([8])
print(col_mean.shape) # torch.Size([5])
print(mx.shape) # torch.Size([8])# keepdim=True면 broadcasting 맞추기 편함!!
X = torch.randn(8, 5)
mu1 = X.mean(dim=1, keepdim=False)
mu2 = X.mean(dim=1, keepdim=True) # (8, 1)
X_norm = X - mu2 # broadcast
# keepdim=False이면 차원 제거, keepdim=True이면 차원 유지
print(mu1.shape) # torch.Size([8])
print(mu2.shape) # torch.Size([8, 1])
print(X_norm.shape) # torch.Size([8, 5])
4. (중요!!) 배치별로 다른 인덱스 뽑기/쓰기 : gather / scatter
# gather: 배치별 정답 label 위치의 logit만 뽑기
B, C = 4, 6
logits = torch.randn(B, C)
print(logits)
labels = torch.tensor([1, 3, 0, 5]) # (B,)
picked = logits.gather(1, labels[:, None]).squeeze(1) # (B,)
print(picked)
'''
tensor([[ 0.8128, -1.5767, 0.6766, 1.6483, 1.5687, -0.6723],
[-0.2185, 1.4773, -0.2952, 0.2396, 1.0336, -2.0708],
[ 0.2267, -0.6256, 0.9033, 0.9330, -1.8699, 0.8545],
[-0.3218, 0.7439, -1.8198, 0.1595, 0.9763, 1.9887]])
tensor([-1.5767, 0.2396, 0.2267, 1.9887])
'''
# scatter_: one-hot 만들기
onehot = torch.zeros(B, C)
out = onehot.scatter_(1, labels[:, None], 1.0) # (B, C)
print(out)
'''
tensor([[0., 1., 0., 0., 0., 0.],
[0., 0., 0., 1., 0., 0.],
[1., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 1.]])
'''
5. 슬라이딩 윈도우/시퀀스 벡터화: unfold
T, F = 100, 9
x = torch.randn(T, F)
window, stride = 16, 4
wins = x.unfold(dimension=0, size=window, step=stride)
print(wins.shape) # torch.Size([22, 9, 16])
6. 배치 행렬 연산: matmul, bmm, einsum
- for문으로 각 배치마다 matmul하던 작업을 한 줄로 수행할 수 있음
B, N, M, K = 32, 10, 20 ,8
A = torch.randn(B, N, K)
Bmat = torch.randn(B, K, M)
Y = torch.bmm(A, Bmat) # (B, N, M)
print(Y.shape) # torch.Size([32, 10, 20])X = torch.randn(4, 5, 16)
Y = torch.randn(4, 7, 16)
dots = torch.einsum('bnd,bmd->bnm', X, Y)
print(dots.shape) # torch.Size([4, 5, 7])
번외) Training 코드에서 실제로 적용해본다면...?
1. 배치별 loss를 for로 누적하여 한 번에 계산
B, C = 100, 6
logits = torch.randn(B, C) # (100, 6)
labels = torch.tensor(torch.arange(B)) # (100,)
# bad
def bad():
loss = 0.
for i in range(B):
loss += torch.nn.functional.cross_entropy(logits[i:i+1], labels[i:i+1])
# good
loss = 0.
def good():
loss = torch.nn.functional.cross_entropy(logits, labels)
2. 조건별 값 변경 : where / masked_fill
# bad: if/for
# good:
x = torch.randn(2, 3)
print(x)
x = x.masked_fill(x < 0, 0)
print(x)
'''
tensor([[ 1.2381, 0.5523, 0.2781],
[-0.8014, 0.0522, -0.3672]])
tensor([[1.2381, 0.5523, 0.2781],
[0.0000, 0.0522, 0.0000]])
'''
3-3. 벡터화 공식 5가지 정리
마지막으로, 벡터화 스킬은 여러 가지가 있어서 따로 정해져있는 공식은 없지만,,,
개인적으로 아래 5가지만이라도 확실하게 암기해두고 있으면 코드에 벡터화를 적용하기에 훨씬 수월할 것 같아서 정리해보았다.
- per-element if → where / mask
- NumPy: np.where, x[mask]=..
- Torch: torch.where, masked_fill, boolean indexing
- per-row / per-col 집계 → sum/mean/max(..., axis/dim=)
- 루프를 돌면서 누적하지 말고 축을 지정
- per-sample 특정 인덱스 값 뽑기 → gather / take_along_axis
- "배치마다 다른 위치"를 뽑는 건 인덱싱 벡터화를 쓰면 됨
- one-hot / 특정 위치에 쓰기 → scatter_
- 루프로 one-hot 만들지 말자!!
- sliding window→ unfold(Torch) / sliding_window_view (NumPy)
- 윈도우 자체를 한 번에 생성
'Python' 카테고리의 다른 글
| [Python] Polars 라이브러리로 대용량 데이터 다루기 | Polars 사용법 및 기본 연산 예제 코드 (0) | 2025.12.28 |
|---|---|
| [Python] DFS BFS 개념 정리 | 깊이 우선 탐색과 너비 우선 탐색 설명 및 예제 코드 (6) | 2025.08.04 |
| [Python] 구현 (Implementation) : 이론 및 예제 문제 풀이 (6) | 2025.07.27 |
| [프로그래머스/Python] 42885 : 구명보트 - 그리디 알고리즘 (3) | 2025.07.27 |
| [프로그래머스/Python] 12982 : 예산 - 그리디 알고리즘 (2) | 2025.07.26 |
