텐서 연산과 모양 다루기

연산 전에 shape부터 확인하기

PyTorch 코드를 읽을 때 가장 먼저 볼 것은 값보다 shape입니다. 딥러닝 모델은 텐서를 입력받고 텐서를 내보냅니다. 이때 모양이 맞지 않으면 계산이 멈춥니다.

import torch

x = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])

print(x.shape)  # torch.Size([2, 3])

[2, 3]은 행이 2개, 열이 3개라는 뜻입니다. 이미지 데이터라면 보통 [배치, 채널, 높이, 너비]처럼 더 많은 차원을 가집니다. 앞으로 코드를 작성할 때는 연산 전후에 shape를 자주 출력하는 습관을 들이면 오류를 훨씬 빨리 찾을 수 있습니다.

인덱싱과 슬라이싱

텐서는 NumPy 배열처럼 원하는 위치를 꺼낼 수 있습니다.

x = torch.tensor([[10, 20, 30],
                  [40, 50, 60]])

print(x[0])      # 첫 번째 행: tensor([10, 20, 30])
print(x[1, 2])   # 두 번째 행, 세 번째 열: tensor(60)
print(x[:, 0])   # 모든 행의 첫 번째 열: tensor([10, 40])
print(x[0, :2])  # 첫 번째 행에서 앞의 두 개: tensor([10, 20])

:는 "전체" 또는 "범위"를 뜻합니다. 예를 들어 이미지 배치가 images.shape == [32, 3, 224, 224]라면 images[:8]은 32장 중 앞의 8장만 꺼냅니다. 이때 첫 번째 차원은 보통 배치 차원입니다.

요소별 연산과 행렬곱은 다르다

같은 크기의 텐서끼리는 사칙연산을 할 수 있습니다.

a = torch.tensor([[1, 2],
                  [3, 4]])
b = torch.tensor([[10, 20],
                  [30, 40]])

print(a + b)
print(a * b)

a * b는 같은 위치의 값끼리 곱하는 요소별 곱셈입니다. 즉 1*10, 2*20처럼 계산합니다.

반면 행렬곱은 행과 열을 조합해 계산합니다. 딥러닝의 선형 계층에서 자주 쓰이는 방식입니다.

print(torch.matmul(a, b))
print(a @ b)

torch.matmul(a, b)와 a @ b는 같은 의미로 쓸 수 있습니다. 행렬곱을 하려면 앞 텐서의 마지막 크기와 뒤 텐서의 뒤에서 두 번째 크기가 맞아야 합니다. 예를 들어 [2, 3] @ [3, 4]는 가능하고 결과는 [2, 4]가 됩니다.

모델 입력에 맞게 모양 바꾸기

모델에 넣기 전에는 텐서 모양을 바꿔야 할 때가 많습니다. 대표적으로 이미지 한 장 28 x 28을 선형 계층에 넣으려면 길이 784의 벡터로 펼쳐야 합니다.

x = torch.randn(32, 1, 28, 28)  # 배치 32개, 채널 1개, 28x28 이미지

flat = x.reshape(32, 784)
print(flat.shape)  # torch.Size([32, 784])

view와 reshape는 모양을 바꿀 때 사용합니다. 초보자는 우선 reshape를 쓰면 됩니다. view는 메모리 배치 조건에 따라 오류가 날 수 있습니다.

x = torch.randn(32, 1, 28, 28)

print(x.flatten(start_dim=1).shape)  # [32, 784]

flatten(start_dim=1)은 첫 번째 차원인 배치는 유지하고, 나머지 차원을 모두 펼칩니다. 모델 입력을 만들 때 자주 쓰는 안전한 방식입니다.

차원을 추가하거나 제거할 때는 unsqueeze와 squeeze를 사용합니다.

x = torch.randn(28, 28)
print(x.unsqueeze(0).shape)  # [1, 28, 28]

x = torch.randn(1, 3, 224, 224)
print(x.squeeze(0).shape)    # [3, 224, 224]

unsqueeze(0)은 맨 앞에 차원 하나를 추가합니다. 이미지 한 장을 배치처럼 다루고 싶을 때 유용합니다. squeeze는 크기가 1인 차원을 제거합니다. 단, 제거하면 안 되는 배치 차원까지 사라질 수 있으니 어느 차원을 지울지 지정하는 편이 안전합니다.

cat과 stack의 차이

여러 텐서를 합칠 때는 torch.cat과 torch.stack을 자주 씁니다. 둘은 비슷해 보이지만 결과 차원이 다릅니다.

a = torch.randn(2, 3)
b = torch.randn(2, 3)

print(torch.cat([a, b], dim=0).shape)    # [4, 3]
print(torch.cat([a, b], dim=1).shape)    # [2, 6]
print(torch.stack([a, b], dim=0).shape)  # [2, 2, 3]

cat은 기존 차원 방향으로 이어 붙입니다. dim=0이면 배치가 늘고, dim=1이면 특성 수가 늘어납니다.

stack은 새 차원을 만들어 쌓습니다. 그래서 [2, 3] 텐서 두 개를 쌓으면 [2, 2, 3]처럼 차원이 하나 늘어납니다. "이어 붙일지", "새 묶음으로 만들지"를 먼저 생각하면 선택이 쉬워집니다.

브로드캐스팅은 편하지만 조심해야 한다

브로드캐스팅은 크기가 다른 텐서끼리도 가능한 경우 자동으로 크기를 맞춰 계산하는 규칙입니다.

x = torch.randn(4, 3)
bias = torch.randn(3)

print((x + bias).shape)  # [4, 3]

bias의 shape는 [3]이지만 PyTorch는 이를 [4, 3]처럼 확장해서 더합니다. 이 기능은 매우 편리합니다. 하지만 의도하지 않은 방향으로 확장되면 결과가 틀릴 수 있습니다.

x = torch.randn(4, 3)
y = torch.randn(4)

# x + y  # 오류: 마지막 차원 3과 4가 맞지 않음

브로드캐스팅은 뒤쪽 차원부터 비교합니다. 크기가 같거나, 둘 중 하나가 1이면 계산할 수 있습니다. 오류가 난다면 두 텐서의 shape를 나란히 출력하고 마지막 차원부터 비교해 보세요.

shape를 읽는 연습이 실력이다

텐서 조작은 외울 것이 많아 보이지만 기준은 단순합니다. 먼저 현재 shape를 확인합니다. 그다음 원하는 결과 shape를 정합니다. 마지막으로 인덱싱, 연산, reshape, flatten, cat, stack 중 필요한 도구를 고릅니다.

모델 코드는 결국 텐서의 모양을 맞춰 가는 과정입니다. 값이 맞아도 모양이 틀리면 학습은 시작되지 않습니다. 다음 장부터 NumPy 변환과 장치 이동까지 배우면, 텐서를 실제 데이터 처리 흐름 안에서 더 자연스럽게 다룰 수 있습니다.