[Pytorch] Broadcasting: 차원이 다를 때 연산

Pytorch를 이용하여 모델을 구현하다 보면 텐서 간의 덧셈, 곱셈 등 연산을 수행할 때 차원이 서로 다른 경우를 자주 만나게 된다. 그런데도 에러 없이 연산이 가능한 이유는 바로 Broadcasting 때문이다. Broadcasting은 Numpy에서 도입된 개념으로, PyTorch에서도 동일하게 적용된다. 이 기능 덕분에 복잡한 차원 조작 없이도 깔끔한 수식 구현이 가능하다.

Broadcasting이란 서로 shape가 다른 두 tensor 간의 연산에서 작은 텐서의 shape을 자동으로 확장하여 큰 텐서와 연산이 가능하도록 맞춰주는 암시적(implicit) 규칙이다. 이는 개발자가 명시적으로 .expand()나 .repeat()를 호출하지 않아도 자동으로 연산되게 도와주는 기능이다. 아래의 조건이 맞으면 두 개의 tensor는 broadcatable 하다.

• 각 tensor는 최소 one dimension이상 가지고 있어야 한다.
• dimension sizes를 iterate 할 때에 dimension이 같거나 둘 중 하나가 1이거나 둘 중 하나가 없어야 한다.

 

PyTorch의 broadcasting은 다음의 규칙에 따라 동작한다.

1. 오른쪽에서부터 차원을 비교한다.
2. 두 차원이 같거나, 둘 중 하나가 1이면 연산 가능하다.
3. 그렇지 않으면 에러가 발생한다.

 

예를 들어 다음 두 텐서가 있다고 하자.

• A.shape = (4, 3, 2)
• B.shape = ( , 1, 2) ← 생략된 부분은 자동으로 맞춰짐

 

두 텐서는 다음처럼 비교된다.

A: (4, 3, 2)
B: (1, 1, 2) ← 자동으로 앞쪽 차원에 1이 추가됨

 

코드를 통해 알아보자.

 

예제 1: 스칼라와 텐서의 연산

import torch

x = torch.tensor([[1, 2], [3, 4]])  # shape: (2, 2)
y = 10                              # scalar

z = x + y  # y가 (2, 2)처럼 확장됨
print(z)
# tensor([[11, 12],
#        [13, 14]])

 

예제 2: (3, 1) vs (1, 4) 텐서의 덧셈

a = torch.randn(3, 1)  # shape: (3, 1)
# tensor([[ 0.9777],
#         [ 0.4755],
#         [-0.3185]])
b = torch.randn(1, 4)  # shape: (1, 4)
# tensor([[-0.1519,  0.3757,  1.1791,  0.3374]])

c = a + b  # shape: (3, 4)
# tensor([[ 0.8258,  1.3533,  2.1567,  1.3150],
#         [ 0.3236,  0.8512,  1.6546,  0.8129],
#         [-0.4704,  0.0572,  0.8606,  0.0189]])

 

---

 

참고 자료

 

https://velog.io/@optjyy/pytorch-broad-casting