Python squeeze 및 unsqueeze 함수 - view와의 차이

PyTorch에서 텐서의 차원을 다룰 때 매우 자주 사용되는 함수가 squeeze()와 unsqueeze()입니다. 이 두 함수는 텐서의 차원을 제거하거나 추가할 때 사용되며, 딥러닝 모델을 구현할 때 입력 형식을 맞추거나 특정 연산을 위해 매우 자주 등장합니다.

이 글에서는 squeeze()와 unsqueeze()의 개념과 사용법을 설명하고, 다른 view 함수와의 차이도 함께 비교합니다.

 

squeeze()란?

squeeze()는 텐서에서 크기가 1인 차원을 제거하는 함수입니다. 예를 들어, 크기가 (1, 3, 1, 5)인 텐서를 squeeze() 하면 (3, 5)로 바뀝니다. 아래 예시와 같이 크기가 1인 차원이 자동으로 제거됩니다.

import torch

x = torch.randn(1, 3, 1, 5)
print(x.shape)  # torch.Size([1, 3, 1, 5])

y = x.squeeze()
print(y.shape)  # torch.Size([3, 5])

 

특정 차원만 제거

dim 인자를 주면 해당 축이 1일 때만 제거합니다.

x = torch.randn(1, 3, 1, 5)

y = x.squeeze(0)  # 첫 번째 차원 제거
print(y.shape)    # torch.Size([3, 1, 5])

z = x.squeeze(2)  # 세 번째 차원 제거
print(z.shape)    # torch.Size([1, 3, 5])

 

단, 지정한 차원의 크기가 1이 아니면 아무 일도 일어나지 않습니다.

 

unsqueeze()란?

unsqueeze()는 텐서에 크기 1인 차원을 추가하는 함수입니다. 예를 들어 (3, 5) 형태의 텐서에 차원을 추가하면 (1, 3, 5) 또는 (3, 1, 5) 등으로 만들 수 있습니다. 차원은 0부터 시작하며, 해당 위치에 1 크기의 차원이 삽입됩니다. 아래와 같이 추가하고 싶은 차원의 인덱스를 입력으로 줍니다.

x = torch.randn(3, 5)
print(x.shape)  # torch.Size([3, 5])

y = x.unsqueeze(0)
print(y.shape)  # torch.Size([1, 3, 5])

z = x.unsqueeze(2)
print(z.shape)  # torch.Size([3, 5, 1])

 

 

딥러닝 모델에서는 입력 텐서의 shape이 매우 중요합니다. 예를 들어:

• 배치 차원이 없으면 모델에 입력할 수 없습니다 → unsqueeze(0) 사용
• 모델 출력에서 불필요한 차원이 생길 수 있습니다 → squeeze() 사용
• 이미지 데이터의 shape을 맞추기 위해 자주 사용됩니다

---

view와의 관계

squeeze()와 unsqueeze()는 각각 차원을 제거하거나 추가하는 역할을 하지만, 내부적으로는 결국 shape을 바꾸는 것과 같습니다. 따라서 view()를 통해 동일한 결과를 구현할 수도 있습니다.

 

squeeze()와 view() 

아래 예제에서는 1 × 3 × 1 텐서를 [3]으로 만든 것입니다. squeeze()는 크기가 1인 차원을 제거하고, view()는 직접 원하는 shape을 지정합니다.

주의: view()를 사용할 경우에는 최종 shape을 **정확히 알고 있어야 하며**, 자동으로 차원을 제거해주지는 않습니다.

import torch

# 1x3x1 텐서
x = torch.randn(1, 3, 1)
print("원본 shape:", x.shape)  # torch.Size([1, 3, 1])

# squeeze 사용
a = x.squeeze()
print("squeeze 결과:", a.shape)  # torch.Size([3])

# view로 동일하게 구현
b = x.view(3)
print("view 결과:", b.shape)  # torch.Size([3])

 

unsqueeze()와 view() 

unsqueeze(0)은 dim=0에 크기 1인 차원을 추가하고, view(1, 3)은 직접 그런 shape으로 변경하는 것입니다.

# 1차원 텐서
x = torch.tensor([1.0, 2.0, 3.0])
print("원본 shape:", x.shape)  # torch.Size([3])

# unsqueeze로 첫 번째 차원 추가 → shape: [1, 3]
a = x.unsqueeze(0)
print("unsqueeze(0) 결과:", a.shape)

# view로 동일하게 구현
b = x.view(1, 3)
print("view 결과:", b.shape)

 

 

unsqueeze(dim)는 어떤 위치에 "1"을 넣겠다는 **의미 중심의 함수**입니다. 반면 view(...)는 오직 "결과 shape"만 고려하는 **모양 중심의 함수**입니다. 따라서 view()는 의도를 명확히 드러내기 어렵고, 가독성 측면에서는 unsqueeze()가 더 명확한 선택이 됩니다.

감사합니다!