squeeze와 unsqueeze

squeeze(): 크기 1인 차원 제거

squeeze() 는 크기가 1인 차원을 모두 제거합니다. 불필요하게 추가된 차원을 정리할 때 사용합니다.

before: (1, 3, 1, 4)
after:  (3, 4)        — 크기 1인 dim 0, dim 2 제거

import torch

t = torch.randn(1, 3, 1, 4)
print(t.shape)             # torch.Size([1, 3, 1, 4])

# 모든 크기 1 차원 제거
squeezed = t.squeeze()
print(squeezed.shape)      # torch.Size([3, 4])

# 특정 차원만 지정해서 제거
print(t.squeeze(0).shape)  # torch.Size([3, 1, 4])  — dim 0만 제거
print(t.squeeze(2).shape)  # torch.Size([1, 3, 4])  — dim 2만 제거
print(t.squeeze(1).shape)  # torch.Size([1, 3, 1, 4])  — dim 1은 크기 3이라 변화 없음

시각적 이해

원본:    [ 1, 3, 1, 4 ]
           ↑        ↑
         크기 1   크기 1

squeeze 후: [ 3, 4 ]

unsqueeze(0): (2, 3) → (1, 2, 3)

변환 전

형태: [2, 3]

변환 후

형태: [1, 2, 3]

메모리 순서 (1D 평탄화)

# 실제로 데이터는 동일하고 shape만 변함
t = torch.tensor([[[1.0, 2.0, 3.0]]])   # shape: (1, 1, 3)
s = t.squeeze()
print(s)          # tensor([1., 2., 3.])
print(s.shape)    # torch.Size([3])
print(t.data_ptr() == s.data_ptr())  # True — 같은 메모리

unsqueeze(): 새 차원 추가

unsqueeze(dim) 은 지정한 위치에 크기 1인 차원을 삽입합니다. 브로드캐스팅 준비, 배치 차원 추가 등에 자주 사용합니다.

t = torch.tensor([1.0, 2.0, 3.0])   # shape: (3,)

print(t.unsqueeze(0).shape)   # torch.Size([1, 3])  — 앞에 차원 추가
print(t.unsqueeze(1).shape)   # torch.Size([3, 1])  — 뒤에 차원 추가
print(t.unsqueeze(-1).shape)  # torch.Size([3, 1])  — 음수 인덱스 사용 가능

차원 삽입 위치 시각화

원본: (3,)      → [3]

unsqueeze(0):   [1, 3]   — dim 0 앞에 삽입
unsqueeze(1):   [3, 1]   — dim 1 앞에 삽입 (끝)
unsqueeze(-1):  [3, 1]   — 마지막 앞에 삽입 (끝과 동일)

t = torch.randn(3, 4)    # shape: (3, 4)

print(t.unsqueeze(0).shape)    # torch.Size([1, 3, 4])
print(t.unsqueeze(1).shape)    # torch.Size([3, 1, 4])
print(t.unsqueeze(2).shape)    # torch.Size([3, 4, 1])
print(t.unsqueeze(-1).shape)   # torch.Size([3, 4, 1])

배치 차원 추가/제거 패턴

딥러닝에서 모델은 보통 배치 입력 (batch, ...) 을 기대합니다. 단일 샘플을 처리할 때 배치 차원을 추가해야 합니다.

# 단일 이미지: (C, H, W)
image = torch.randn(3, 28, 28)
print(image.shape)   # torch.Size([3, 28, 28])

# 모델 입력을 위해 배치 차원 추가: (1, C, H, W)
batch_image = image.unsqueeze(0)
print(batch_image.shape)   # torch.Size([1, 3, 28, 28])

# 모델 출력 후 배치 차원 제거
output = model(batch_image)          # shape: (1, num_classes)
prediction = output.squeeze(0)       # shape: (num_classes,)

브로드캐스팅을 위한 차원 정렬

# 열벡터와 행벡터의 외적 계산
col = torch.tensor([1.0, 2.0, 3.0])   # shape: (3,)
row = torch.tensor([10.0, 20.0])       # shape: (2,)

# unsqueeze로 차원 정렬 후 브로드캐스팅
outer = col.unsqueeze(1) * row.unsqueeze(0)
# (3, 1) * (1, 2) → (3, 2)
print(outer)
# tensor([[10., 20.],
#         [20., 40.],
#         [30., 60.]])

# 또는 torch.outer() 사용
print(torch.outer(col, row))   # 동일한 결과

squeeze와 unsqueeze의 관계

두 함수는 서로 역연산입니다.

t = torch.randn(3, 4)     # shape: (3, 4)

# unsqueeze 후 squeeze
t_expanded = t.unsqueeze(0)      # shape: (1, 3, 4)
t_restored = t_expanded.squeeze(0)  # shape: (3, 4)
print(torch.equal(t, t_restored))   # True

None 인덱싱 단축 문법

Python의 None (= np.newaxis) 을 인덱싱에 사용하면 unsqueeze 와 동일합니다.

t = torch.randn(3, 4)

# 아래 세 가지는 모두 동일
print(t.unsqueeze(0).shape)    # torch.Size([1, 3, 4])
print(t[None].shape)           # torch.Size([1, 3, 4])
print(t[None, :, :].shape)     # torch.Size([1, 3, 4])

# 중간에 삽입
print(t[:, None, :].shape)     # torch.Size([3, 1, 4])
print(t[:, :, None].shape)     # torch.Size([3, 4, 1])