Pytorch Broadcasting

Broadcastig은 두 텐서가 서로 다른 크기일 때, 자동으로 차원을 맞춰 연산을 수행할 수 있게 해주는 기능이다.이 같은 기능은 반복문을 굳이 작성하지 않아도 내부적으로 알아서 원소별(element-wise) 연산이 가능하게 해준다.

Concept of Broadcasting

Broadcasting은 두 tensor의 차원이 다를 때 작동한다. 작은 tensor의 차원을 큰 tensor와 맞추기 위해 자동으로 확장하여 element-wise 연산을 가능하게 해준다.

1. 뒤에서부터 차원을 맞추기 시작한다.
2. 각 차원에서 둘 중 하나가 1인 경우 해당 차원이 큰 tensor에 맞춰 확장한다.
3. 두 tensor가 서로 다른 크기의 차원을 가지며, 어느 한쪽도 1이 아니라면 Broadcasting은 불가능하다.

import torch

# (3, 1) 모양의 텐서
a = torch.tensor([[1], [2], [3]])  # shape: (3, 1)

# (1, 4) 모양의 텐서
b = torch.tensor([10, 20, 30, 40])  # shape: (1, 4)

# 브로드캐스팅을 통해 원소별 덧셈
result = a + b
print(result)

# answer
tensor([[11, 21, 31, 41],
        [12, 22, 32, 42],
        [13, 23, 33, 43]])

즉, a의 tensor의 크기를 (3, 4)로 b의 tensor의 크기를 (3, 4)로 확장한다. 따라서 최종적으로 확장된 두 텐서를 element-wise로 연산하게 된다.

yolov3에 나오는 코드이긴 한데,

iou_anchors = iou(torch.tensor(box[2:4], self.anchors)

이 코드를 보면 box[2:4]의 shape는 (2,)이고 self.anchors의 shape는 (9, 2)이다. Broadcasting을 생각해보면 box[2:4]의 shape는 (9, 2)로 확장된다.

예를들어 box[2:4]가 [0.11, 0.33] 이라면

[
[0.11, 0.33],
[0.11, 0.33],
[0.11, 0.33],
[0.11, 0.33],
[0.11, 0.33],
[0.11, 0.33],
[0.11, 0.33],
[0.11, 0.33],
[0.11, 0.33],
]

로 확장되어 계산되는 것이다.

계산 속도와 메모리 효율성을 생각하면 for보다 Broadcasting을 사용하는 방식이 훨씬 좋다. Broadcasting은 벡터화된 연산을 활용하므로 반복문을 사용하지 않아도 텐서 연산을 빠르게 수행할 수 있기 때문이다.