[PYTORCH] 차원을 자유자재로 다루는 2가지 방법: squeeze()와 unsqueeze() 완벽 해결 가이드

squeeze()와 unsqueeze()

 

딥러닝 모델을 설계하다 보면 가장 빈번하게 마주치는 난관 중 하나가 바로 'Tensor Shape Mismatch' 에러입니다. 데이터의 본질은 변하지 않았음에도 불구하고, 연산을 수행하는 레이어나 함수가 요구하는 차원(Dimension)이 맞지 않아 코드 실행이 중단되는 경험은 모든 개발자가 겪는 숙명과도 같습니다. 이때 우리는 텐서의 차원을 늘리거나 줄여야 하는 상황에 직면합니다. PyTorch는 이를 위해 squeeze()와 unsqueeze()라는 직관적이고 강력한 도구를 제공합니다. 본 포스팅에서는 단순한 API 설명을 넘어, 실무 딥러닝 파이프라인에서 이 두 함수가 왜 필수적인지, 그리고 실제 모델 배포 환경에서 성능 최적화와 코드 가독성을 위해 어떻게 활용해야 하는지 전문적인 시각에서 심도 있게 다뤄보겠습니다.

---

1. 텐서 차원 조작의 핵심 개념

PyTorch에서 텐서는 다차원 배열입니다. 예를 들어 이미지 데이터는 일반적으로 (Batch, Channel, Height, Width)의 4차원 구조를 가집니다. 하지만 특정 연산(예: Loss Function 계산이나 Linear Layer 통과)에서는 1차원으로 줄어든 결과값이 나오거나, 단일 이미지를 모델에 넣을 때 Batch 차원이 빠진 3차원 상태인 경우가 많습니다. 이때 squeeze와 unsqueeze는 데이터의 총 요소 수(Element count)를 유지하면서 1인 차원만을 제거하거나 추가하는 역할을 수행합니다.

왜 View나 Reshape 대신 이들을 쓰는가?

view()나 reshape()은 전체적인 형태를 완전히 바꿀 수 있는 범용적인 함수입니다. 반면 squeeze()와 unsqueeze()는 오직 '사이즈가 1인 차원'만을 타겟팅합니다. 이는 의도치 않은 데이터 변형을 방지하고 코드의 의도를 명확히 전달(Semantic clarity)하는 데 큰 장점이 있습니다.

---

2. squeeze()와 unsqueeze() 기능 비교 분석

두 함수의 차이점을 한눈에 파악할 수 있도록 비교표로 정리했습니다.

기능 구분	squeeze()	unsqueeze()
핵심 목적	크기가 1인 차원 제거 (축소)	지정한 위치에 크기가 1인 차원 삽입 (확장)
데이터 변화	차원 수가 감소하거나 동일	차원 수가 반드시 1 증가
주요 파라미터	dim (선택 사항)	dim (필수 사항)
주요 사용 사례	연산 결과 불필요한 1차원 정리	싱글 데이터를 배치 데이터로 변환
메모리 참조	원본 텐서의 뷰(View) 반환	원본 텐서의 뷰(View) 반환

---

3. 실무 적용을 위한 핵심 Sample Example 7선

개발자가 실제 프로젝트 현장에서 바로 복사하여 활용할 수 있는 실전 예제들입니다.

Example 1: 단일 이미지의 배치 차원 추가 (Inference 단계)

import torch

# 모델은 (B, C, H, W)를 기대하지만, 불러온 이미지는 (3, 224, 224)
single_img = torch.randn(3, 224, 224)

# 0번째 인덱스에 배치 차원 추가 -> (1, 3, 224, 224)
input_tensor = single_img.unsqueeze(dim=0)
print(f"Input Shape: {input_tensor.shape}")

Example 2: 회귀 분석 결과의 불필요한 차원 제거

# 모델 출력 값이 (Batch, 1) 형태로 나올 때, 스칼라 리스트로 변환하기 전
output = torch.randn(32, 1)

# 마지막 차원이 1인 경우 제거 -> (32,)
flat_output = output.squeeze()
print(f"Squeezed Shape: {flat_output.shape}")

Example 3: 특정 위치에 차원 삽입하여 브로드캐스팅(Broadcasting) 유도

# 텐서 A: (10,), 텐서 B: (10, 5)
# A를 B의 각 행에 곱하고 싶을 때
A = torch.randn(10)
B = torch.randn(10, 5)

# A를 (10, 1)로 만들어 브로드캐스팅 적용
result = A.unsqueeze(1) * B
print(f"Result Shape: {result.shape}")

Example 4: 다중 차원 squeeze (Safety 조작)

# (1, 3, 1, 128) 처럼 여러 개의 1이 섞여 있을 때
complex_tensor = torch.randn(1, 3, 1, 128)

# 특정 차원만 골라서 제거 (0번 차원만 제거)
safe_squeeze = complex_tensor.squeeze(dim=0) # (3, 1, 128)
# 모든 1 제거
full_squeeze = complex_tensor.squeeze() # (3, 128)

Example 5: Grayscale 이미지의 채널 차원 생성

# (H, W) 형태의 그레이스케일 데이터를 CNN 입력용 (1, C, H, W)로 변환
gray_img = torch.randn(28, 28)

# 차례대로 0번(Batch), 1번(Channel) 추가
input_ready = gray_img.unsqueeze(0).unsqueeze(0)
print(f"CNN Input Shape: {input_ready.shape}") # (1, 1, 28, 28)

Example 6: 시계열 데이터(RNN/LSTM) 레이어 맞춤

# (Sequence_Length, Batch) 데이터를 (Seq, Batch, Feature)로 확장
seq_data = torch.randn(10, 32)
# Feature 차원(마지막)에 1 추가
rnn_input = seq_data.unsqueeze(-1)
print(f"RNN Input Shape: {rnn_input.shape}") # (10, 32, 1)

Example 7: Loss 계산을 위한 Label 차원 맞추기

# CrossEntropyLoss 등에서 Target이 (N,)이어야 하는데 (N, 1)로 들어오는 경우
target = torch.empty(64, 1, dtype=torch.long).random_(10)
target = target.squeeze(1) # (64,)
print(f"Target Shape for Loss: {target.shape}")

---

4. 고수들의 팁: 성능과 안정성을 생각하는 해결 방법

단순히 함수를 호출하는 것을 넘어, 대규모 모델을 다룰 때 주의해야 할 3가지 포인트를 제안합니다.

1. 명시적 차원 지정: squeeze()를 인자 없이 사용하면 의도치 않게 중간에 있는 '사이즈 1'인 데이터 차원까지 사라질 수 있습니다. squeeze(dim=1)처럼 항상 명시적으로 차원을 지정하는 습관이 버그를 줄입니다.
2. Negative Index 활용: unsqueeze(-1)은 마지막 차원에 새로운 축을 추가하겠다는 뜻입니다. 텐서의 전체 차원 개수를 매번 계산할 필요가 없어 매우 효율적입니다.
3. In-place 연산 주의: squeeze_()와 같이 언더바(_)가 붙은 메서드는 원본을 직접 수정합니다. 메모리 절약에는 유리하지만, Gradient 전파 단계에서 문제를 일으킬 수 있으므로 autograd 범위 내에서는 가급적 squeeze()를 사용하는 것이 안전합니다.

---

5. 결론

PyTorch의 squeeze()와 unsqueeze()는 데이터의 기하학적 구조를 다루는 마법과 같습니다. Squeeze는 "불필요한 군더더기를 꽉 짜서 없애는 것"이고, Unsqueeze는 "데이터가 들어갈 빈 방을 하나 만드는 것"이라고 이해하면 쉽습니다. 이 두 함수만 능숙하게 다뤄도 딥러닝 구현 시 발생하는 에러의 70% 이상을 사전에 예방할 수 있습니다.

참고 및 출처:
- PyTorch 공식 문서 (torch.squeeze, torch.unsqueeze)
- Deep Learning with PyTorch (Eli Stevens et al.)