차원(Dimension)을 추가하는 가장 쉬운 방법
PyTorch를 공부하다 보면 모델에 데이터를 입력하기 전에 unsqueeze()를 사용하는 코드를 자주 볼 수 있습니다.
x = x.unsqueeze(0)또는
torch.unsqueeze(x, 1)처음에는 “왜 갑자기 차원을 하나 추가하지?”라는 의문이 들기 쉽습니다.
이번 글에서는 torch.unsqueeze()가 무엇인지, 왜 필요한지, 그리고 실제 딥러닝에서 어떻게 사용하는지 예제를 통해 쉽게 알아보겠습니다.
torch.unsqueeze()란?
torch.unsqueeze()는 텐서(Tensor)에 크기가 1인 새로운 차원(Dimension)을 추가하는 함수입니다.
기존 데이터는 그대로 유지하면서 축(axis) 하나만 추가합니다.
기본 형태는 다음과 같습니다.
torch.unsqueeze(input, dim)또는
input.unsqueeze(dim)input- 입력 Tensor
dim- 새 차원을 추가할 위치
예제 1
먼저 1차원 Tensor를 하나 만들어보겠습니다.
import torch
x = torch.tensor([1, 2, 3, 4])
print(x)출력
tensor([1, 2, 3, 4])shape를 확인하면
print(x.shape)torch.Size([4])즉,
(4,)인 1차원 Tensor입니다.
dim=0에 차원 추가
이번에는 앞쪽에 차원을 추가해보겠습니다.
print(torch.unsqueeze(x, 0))결과
tensor([[1, 2, 3, 4]])shape
print(torch.unsqueeze(x, 0).shape)torch.Size([1, 4])기존
[1 2 3 4]↓
[[1 2 3 4]]즉
(4,)↓
(1,4)이 됩니다.
행(Row)이 하나 생긴 것처럼 볼 수 있습니다.
dim=1에 차원 추가
이번에는 뒤쪽에 차원을 추가합니다.
print(torch.unsqueeze(x, 1))결과
tensor([[1],
[2],
[3],
[4]])shape
print(torch.unsqueeze(x,1).shape)torch.Size([4,1])기존
[1 2 3 4]↓
[[1]
[2]
[3]
[4]]즉
(4,)↓
(4,1)열(Column)이 하나 생긴 것처럼 볼 수 있습니다.
dim에 따라 결과가 달라지는 이유
같은 데이터를 사용했지만 결과는 완전히 달라졌습니다.
| 코드 | Shape |
|---|---|
x | (4,) |
x.unsqueeze(0) | (1,4) |
x.unsqueeze(1) | (4,1) |
즉,
dim=0- 가장 앞에 차원을 추가
dim=1- 두 번째 위치에 차원을 추가
입니다.
실행결과
원래 Tensor
unsqueeze(0)
unsqueeze(1)
왜 필요한가?
딥러닝 모델은 입력 데이터의 차원 수가 정확해야 합니다.
예를 들어,
이미지 분류 모델은 보통
(batch, channel, height, width)형태를 기대합니다.
이미지가 한 장이라면
(3,224,224)이지만,
모델은
(1,3,224,224)처럼 Batch 차원까지 포함된 입력을 요구하는 경우가 많습니다.
이때 사용하는 것이 바로
image = image.unsqueeze(0)입니다.
Batch 차원을 하나 추가하는 것입니다.
또 다른 예
벡터를 행렬로 바꾸고 싶을 때도 자주 사용합니다.
x = torch.tensor([1,2,3])
print(x.unsqueeze(0))결과
[[1 2 3]]반대로
print(x.unsqueeze(1))결과
[[1]
[2]
[3]]이렇게 행 벡터(Row Vector)와 열 벡터(Column Vector)를 쉽게 만들 수 있습니다.
squeeze()와의 관계
unsqueeze()의 반대 함수는 squeeze()입니다.
예를 들어
x = torch.tensor([[1,2,3,4]])shape
(1,4)여기서
x.squeeze()결과
tensor([1,2,3,4])shape
(4,)즉,
unsqueeze()- 크기가 1인 차원을 추가
squeeze()- 크기가 1인 차원을 제거
입니다.
실제 딥러닝에서 가장 많이 사용하는 예
예를 들어 이미지 한 장을 모델에 입력한다고 가정해봅니다.
image = torch.randn(3,224,224)
print(image.shape)torch.Size([3,224,224])모델은
(batch, channel, height, width)를 요구하므로
image = image.unsqueeze(0)
print(image.shape)결과
torch.Size([1,3,224,224])이제 모델에 바로 입력할 수 있습니다.
정리
torch.unsqueeze()는 데이터를 변경하지 않고 크기가 1인 새로운 차원을 추가하는 함수입니다.
핵심 내용은 다음과 같습니다.
unsqueeze(dim)은 지정한 위치에 새로운 차원을 추가한다.unsqueeze(0)은 가장 앞에 차원을 추가하여(1, N)형태를 만든다.unsqueeze(1)은 두 번째 위치에 차원을 추가하여(N, 1)형태를 만든다.- 이미지나 시계열 데이터를 모델에 입력할 때 Batch 차원을 추가하기 위해 매우 자주 사용된다.
squeeze()는unsqueeze()의 반대 기능으로, 크기가 1인 차원을 제거한다.
처음에는 단순히 “차원을 하나 늘리는 함수”처럼 보이지만, 실제로는 딥러닝 모델의 입력 형식을 맞추기 위해 가장 자주 사용하는 Tensor 조작 함수 중 하나입니다. unsqueeze()의 동작 원리를 이해하면 Tensor의 Shape를 다루는 감각이 훨씬 자연스러워질 것입니다.