Tensor 연산에서 dimension을 지정하는 것에 관하여

텐서 연산을 시각화한 이미지 created by DALL・E

 

numpy와 pytorch를 공부하면 자주 맞딱뜨리면서도 매번 헷갈리는 개념이 있습니다.

 

바로 행렬 연산에서 dimension을 지정하는 것인데요.

 

다음 예시 행렬을 보겠습니다:

a = torch.arange(0, 18).view(3, 6)
a
>>> tensor([[ 0,  1,  2,  3,  4,  5],
            [ 6,  7,  8,  9, 10, 11],
            [12, 13, 14, 15, 16, 17]])

a는 $3 \times 6$ 행렬 또는 텐서입니다.

 

우리는 일반적으로 dim=0은 행을, dim=1을 열로 알고있고 실제로도 그렇습니다.

 

그렇다면, 위의 예시 행렬에서 행의 값들을 더해서 3개의 값을 반환 받고싶다면 어떻게 해야할까요?

 

행의 값들을 더하고 싶기 때문에 a.sum(dim = 0)을 실행하면 될까요?

 

한 번 실행해보겠습니다

a.sum(dim = 0)
>>> tensor([18, 21, 24, 27, 30, 33])

 

오히려 각 열의 값들을 더한 결과가 나왔습니다

 

반대로 dim=1로 두고 계산해보겠습니다

a.sum(dim=1)
>>> tensor([15, 51, 87])

 

이번엔 앞에서 우리가 원했던 행의 값들을 더한 결과가 벡터로 반환되었습니다

 

그렇다면 이런 현상을 어떻게 이해하면 좋을까요?

 

저는 이것을 특정 dimension을 차곡차곡 포개는 것으로 연상했습니다

 

위 a 텐서를 예시로 설명해보자면

a
>>> tensor([[ 0,  1,  2,  3,  4,  5],
    	    [ 6,  7,  8,  9, 10, 11],
       	    [12, 13, 14, 15, 16, 17]])

 

dim=0으로 지정해주었을 경우

tensor([ 0,  1,  2,  3,  4,  5])
	
    		+
        
tensor([ 6,  7,  8,  9, 10, 11])
		
        	+
        
tensor([12, 13, 14, 15, 16, 17])

이런 방식으로 각 행(dim=0)을 따로 떼어내어 세 개의 서로 다른 텐서로 본 후, 각 텐서를 더해주는 것으로 이해할 수 있습니다.

 

마찬가지로, dim=1로 지정해주었을 경우

tensor([[0],	  tensor([[1],	  tensor([[2],	  tensor([[3],   tensor([4],    tensor([[5],
	[6],	+	  [7],	+	  [8],	+	  [9],  +       [10], +         [11],
        [12]])		  [13]])	  [14]])          [15]])	[16]])          [17]])

이런 방식으로 각 열(dim=1)을 기준으로 분리하여 서로 다른 텐서 6개로 만든 후 각 텐서를 더해주는 것으로 이해할 수 있습니다

---

 

이번에는 3차원 텐서에 이 방법을 적용해보겠습니다 (차원이 높아짐에 따라 dim=0은 깊이를, dim=1이 행을, dim=2가 열을 의미하게 됩니다)

 

예시 행렬 a를 변형하여 $2 \times 3 \times 3$ 텐서를 만들었습니다

a
>>> tensor([[[ 0,  1,  2],
             [ 3,  4,  5],
             [ 6,  7,  8]],

            [[ 9, 10, 11],
             [12, 13, 14],
             [15, 16, 17]]])

(여기서 부터는 tensor라는 표현을 제외하고 수와 괄호만으로 표현하도록 하겠습니다)

 

dim=0을 기준으로 sum연산을 할 경우

[[ 0,  1,  2],
 [ 3,  4,  5],
 [ 6,  7,  8]]
      +
[[ 9, 10, 11],
 [12, 13, 14],
 [15, 16, 17]]
      ||
[[ 9, 11, 13],
 [15, 17, 19],
 [21, 23, 25]]

와 같이 연산이 되어 최종적으로 $3\times 3$ 텐서가 나왔습니다.

 

여기서 잠깐. $2\times 3\times 3$ 텐서에서 dim=0을 기준으로 sum 연산을 했더니 $3\times 3$ 텐서가 나왔습니다. 즉, 우리가 지정한 dim=0의 차원이 제거되는 효과가 생겼습니다.

 

그렇다면 dim=1, dim=2로 지정한 경우에는 sum 연산 후 출력 텐서가 어떻게 될까요?

 

이 경우에는 dim=1, dim=2의 차원 수가 3으로 같으므로 둘 모두 출력 결과가 $2\times 3$ 텐서가 생성될 것입니다.

 

dim=1로 지정하여 연산하는 과정과 결과를 보겠습니다:

[[0,  1,  2]
      +
 [3,  4,  5]
      +
 [6,  7,  8]]

[[9,  10, 11]
      +
 [12, 13, 14]
      +
 [15, 16, 17]]
     
     ||
[[9, 12, 15],
 [36, 39, 42]]

이번에도 각 행을 분리하여 합치는 연산을 하여 나온 두 결과벡터를 이어 붙여 $2\times 3$ 텐서가 생성된 것을 확인할 수 있습니다

 

dim=2의 경우도 한 번 직접 풀어보세요! 정답은 아래에 적어두겠습니다.

a.sum(dim=2)
>>> tensor([[[ 0,  1,  2],
             [ 3,  4,  5],
             [ 6,  7,  8]],

            [[ 9, 10, 11],
             [12, 13, 14],
             [15, 16, 17]]])

---