np.squeeze, np.expand_dims(unsqueeze) 이해

• 쓸모없는 차원 1을 제거(squeeze) 또는 추가(expand_dims)
• expand_dim의 경우, pytorch의 unsqueeze와 같다.

예시

 

sample_array = np.array([[1,2,3], [4,5,6]])

# axis를 기준으로 차원 추가
sample_array_expand_0 = np.expand_dims(sample_array, axis=0)
sample_array_expand_1 = np.expand_dims(sample_array, axis=1)
sample_array_expand_2 = np.expand_dims(sample_array, axis=2)

# axis를 기준으로 차원 제거(차원이 1인경우만 제거 가능)
sample_array_expand_1_squeeze = np.squeeze(sample_array_expand_1, axis=1)

print(sample_array.shape)
# (2, 3)

print(sample_array_expand_0.shape)
# (1, 2, 3)
print(sample_array_expand_1.shape)
# (2, 1, 3)
print(sample_array_expand_2.shape)
# (2, 3, 1)

print(sample_array_expand_1_squeeze.shape)
# (2, 3)