재미있는 삽질 기록

1. dim 쉽게 이해하기

#pytorch#tensor#dim

한참 Pytorch를 공부하고 있을 때, tensor.sum(dim=x), torch.stack((tensors), dim=x) 등의 연산에서 dim=x가 어떤 방향으로 연산이 이루어지는 건지 이해하기 어려웠었는데, 쉽게 이해할 수 있었던 방법을 정리해봤다.

단순한 규칙 알아두기#

  1. 연산은 dim=x의 방향으로 이루어진다.
  2. 연산되는 tensor의 shape를 구한다.
  3. shape에서 dim=x를 제거하거나 추가한 값이 연산 후 텐서의 shape가 된다.

핵심은 3번이다. 어느 방향으로 화살표가 가는가를 외우려고 하면 고차원으로 갈 수록 더 헷갈리는데, 연산의 종류에 따라서 shape에서 그 자리가 사라지거나, 생긴다로 보면 한 번에 정리된다.

차원 축소 연산: dim 삭제#

x = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])   # shape: (2, 3)

x.sum(dim=0)   # tensor([5, 7, 9])  -> shape: (3,)
x.sum(dim=1)   # tensor([6, 15])    -> shape: (2,)
  • dim=0으로 더하면 shape (2, 3)에서 0번 자리가 사라져 (3,)이 된다. 즉 행을 따라 위아래로 더해져 열별 합이 나온다.
  • dim=1로 더하면 1번 자리가 사라져 (2,)가 된다. 즉 각 행 안에서 더해져 행별 합이 나온다.

방향을 외우는 대신 dim=0을 없앴으니 (3,)이 남는다로 생각하면 결과 모양이 바로 보인다.

차원 확장 연산: 새로운 dim 추가#

a = torch.tensor([1, 2, 3])   # shape: (3,)
b = torch.tensor([4, 5, 6])   # shape: (3,)

torch.stack((a, b), dim=0)    # shape: (2, 3)
torch.stack((a, b), dim=1)    # shape: (3, 2)
  • dim=0으로 쌓으면 0번 자리에 새 축이 생겨 (2, 3)
  • dim=1로 쌓으면 1번 자리에 새 축이 생겨 (3, 2)

sum이 제거, 즉 차원을 축소시키는 연산이라면 stack은 차원을 확장하는 메서드이다. 차원을 제거하냐 확장하느냐의 문제이지, 규칙은 동일하다.

stack 외에 unsqueeze 도 같은 부류다. 반대로 cat은 새 축을 만들지 않고 기존 축을 늘리기만 해서 차원 수는 그대로라는 점을 주의하자.

dim=0는 가장 최근에 생겨난 차원#

dim 자체를 바라보는 관점도 하나 잡아두면 좋다. 강의를 들으면서 메모해둔 게 있는데, 항상 dim=0이 어디인지부터 찾자는 것이다. dim은 col, row 같은 고정된 개념이 아니다. dim은 가장 최근에 생겨난 차원의 방향이다.

예를 들어 shape가 [3, 2, 4]인 3-D Tensor가 있다고 하자. 안쪽의 (2, 4) 행렬 M이 3개 쌓여 있는 구조다. 여기서 dim=0은 그 M들이 쌓인 방향, 즉 가장 바깥쪽이자 가장 마지막에 더해진 차원이다.

  • [3, 2, 4] 라면 (2, 4) 행렬 M이 3개
  • [10, 3, 2, 4] 라면 그 3개짜리 묶음이 다시 10개

이렇게 새 차원이 앞쪽(dim=0)에 붙는다고 보면, 차원이 늘어날 때 dim 번호가 어떻게 밀리는지가 헷갈리지 않는다.

channel first vs channel last#

이 관점이 실제로 쓰이는 대표적인 예가 이미지의 채널 순서다.

  • channel first (C, H, W): PyTorch가 대표적이다. 배치까지 포함하면 NCHW이고, Caffe·cuDNN도 이 방식을 쓴다.
  • channel last (H, W, C): TensorFlow/Keras의 기본 포맷이다(NHWC). 이미지를 다루는 matplotlib, PIL, NumPy도 채널이 마지막에 온다.

그래서 같은 이미지라도 프레임워크마다 채널 위치가 달라진다. torch Tensor를 matplotlib으로 그리려면 채널 위치를 옮겨줘야 하는데, 이때 permute로 channel first를 channel last로 바꾼다.

img = torch.rand(3, 28, 28)   # (C, H, W) - torch 기본
img.permute(1, 2, 0).shape    # (28, 28, 3) - (H, W, C), plt가 원하는 형태

permute의 인자는 새로 배치할 dim의 순서다. (1, 2, 0)은 기존 dim 1, 2, 0을 차례대로 가져와 (H, W, C)를 만든다.

댓글