[Numpy] np.stack에 대한 이해

np.stack은 쌓고자 하는 array와 어떤 방향으로 쌓을지 결정해주는 axis 파트가 있습니다.

Staking된 결과를 보면 오히려 헷갈리기 때문에, 저 같은 경우는 shape을 기준으로 어떻게 쌓이는지 우선 관찰하는 편입니다.

 

---

 

 

예를 들어 아래의 코드처럼 보면 좀 더 이해가 쉽다는 것입니다.

import numpy as np

a = np.arange(1, 3, 1).reshape(1, 2)
b = np.arange(3, 5, 1).reshape(1, 2)
c = np.arange(5, 7, 1).reshape(1, 2)

s0 = np.stack((a, b, c), axis=0)
s1 = np.stack((a, b, c), axis=1)
s2 = np.stack((a, b, c), axis=2)


print(s0.shape)
# (3, 1, 2)

print(s0)
# [[[1 2]]
#
#  [[3 4]]
#
#  [[5 6]]]

print(s0[0, :, :], s0[1, :, :], s0[2, :, :])
# [[1, 2]], [[3, 4]], [[5, 6]]

print(s1.shape)
# (1, 3, 2)

print(s1)
# [[[1 2]
#   [3 4]
#   [5 6]]]

print(s1[:, 0, :], s1[:, 1, :], s1[:, 2, :])
# [[1, 2]], [[3, 4]], [[5, 6]]

print(s2.shape)
# (1, 2, 3)

print(s2)
# [[[1 3 5]
#   [2 4 6]]]

print(s2[:, :, 0], s2[:, :, 1], s2[:, :, 2])
# [[1, 2]], [[3, 4]], [[5, 6]]

 

shape이 (1, 2)인 array 3개를 axis=0에 쌓는다는 것은, 새로운 좌표의 0번째 좌표에 (1, 2)짜리 array를 세 개 쌓겠다는 뜻이므로 shape이 (3, 1, 2)가 되고, 나머지도 비슷하게 동작합니다.

 

어쨌든 쌓는 재료는 같기 때문에, np.stack에서 지정한 새로운 axis를 따라 slicing해주면, 당연히 재료로 썼던 array들이 다시 등장하게 됩니다.

 

해당 예는 3-dim example이므로 쌓이는 과정을 한 눈에 상상할 수도 있지만, dimension이 그 이상이 되면 당연히 한 눈에 관찰하는 것은 불가능하기 때문에, 지정한 axis 기준으로 쌓이는 것을 위주로 부분적으로 관찰하면 어떻게 쌓이는지 감을 잡을 수 있을 것입니다.