[PYTORCH] torch.cat()과 torch.stack()의 결정적 차이 1가지와 실무 해결 방법 7선

torch.cat()과 torch.stack()

 

 

딥러닝 모델, 특히 파이토치(PyTorch)를 이용해 복잡한 네트워크를 설계하다 보면, 여러 소스에서 생성된 텐서(tensor)를 하나로 합쳐야 하는 상황을 끊임없이 마주하게 됩니다. 예를 들어, RNN이나 Transformer의 각 타임스텝에서 나온 hidden state들을 하나로 모으거나, 여러 이미지 배치를 합쳐 더 큰 배치를 만드는 경우입니다. 이때 우리는 torch.cat()과 torch.stack()이라는 두 가지 강력한 도구를 사용합니다. 하지만 많은 개발자가 이 두 함수의 동작 방식을 명확히 구분하지 못해 의도치 않은 형태(shape)의 텐서를 생성하고, 이는 곧 'dimension mismatch'라는 흔하지만 치명적인 런타임 에러로 이어집니다. 본 포스팅에서는 단순한 정의를 넘어, 이 두 함수의 핵심적인 차이 1가지를 시각적이고 전문적으로 분석하고, 실무에서 마주하는 7가지 구체적인 시나리오에 대한 명쾌한 해결 방법을 코드 예제와 함께 제공합니다. 이 글을 끝까지 읽으시면, 더 이상 텐서 결합 시 형태 mismatch로 고민하지 않게 될 것입니다.

---

1. 서론: 텐서 결합의 난제, 차원(Dimension)

파이토치에서 모든 데이터는 차원의 집합인 텐서로 표현됩니다. 텐서를 결합한다는 것은 단순히 데이터를 옆에 붙이는 것을 넘어, 그 데이터가 위치할 차원의 구조를 결정하는 일입니다. torch.cat()과 torch.stack()은 모두 텐서 리스트를 입력으로 받아 하나의 텐서로 합치지만, 그 데이터를 배치하는 방식에 있어 근본적인 차이가 있습니다.

전문가 팁: 텐서 연산에서 '형태(Shape)'는 데이터의 의미를 결정합니다. 결합 함수를 잘못 선택하면, 데이터의 의미가 완전히 왜곡되거나 연산 자체가 불가능해집니다.

---

2. 결정적 차이: 기존 차원 vs. 새로운 차원

두 함수의 결정적 차이는 '결합이 일어나는 차원이 기존 차원이냐, 아니면 새로 생성된 차원이냐'에 있습니다.

2.1 torch.cat() - Concatenate (기존 차원 활용)

torch.cat()(concatenate)는 입력된 텐서들을 기존에 존재하는 특정 차원(dimension)을 따라 연결합니다. 이 경우, 결합되는 차원을 제외한 다른 차원의 크기는 모두 동일해야 합니다. 결과 텐서의 차원 수(rank)는 입력 텐서들의 차원 수와 동일하게 유지됩니다.

• 동작 방식: 데이터를 기존 차원에 '옆으로' 혹은 '아래로' 계속 이어 붙입니다.
• 조건: 결합하는 차원을 제외한 모든 차원의 크기가 같아야 함.
• 결과: 결합된 차원의 크기가 입력 텐서들의 해당 차원 크기 합이 됨. 차원 수는 변하지 않음.

2.2 torch.stack() - Stack (새로운 차원 생성)

torch.stack()은 입력된 텐서들을 완전히 새로운 차원(new dimension)을 생성하여 그 차원에 쌓습니다. 이 경우, 입력된 모든 텐서들의 형태(shape)가 완벽하게 동일해야 합니다. 결과 텐서의 차원 수(rank)는 입력 텐서들의 차원 수보다 1이 증가합니다.

• 동작 방식: 데이터를 새로운 차원에 '차곡차곡' 쌓습니다. 마치 카드를 쌓는 것과 같습니다.
• 조건: 입력된 모든 텐서들의 형태가 완전히 같아야 함.
• 결과: 지정한 위치에 새로운 차원이 생성되며, 그 차원의 크기는 입력된 텐서의 개수가 됨. 차원 수가 1 증가함.

---

3. 요약 및 비교표

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

비교 항목	torch.cat() (Concatenate)	torch.stack() (Stack)
핵심 동작	기존 차원을 따라 연결	새로운 차원을 생성하여 쌓음
차원 수 (Rank) 변화	변화 없음	1 증가
입력 텐서 조건	결합하는 차원을 제외한 다른 차원의 크기가 동일해야 함	입력된 모든 텐서들의 형태(shape)가 완전히 동일해야 함
결과 텐서 형태 (입력 shape가 (A, B)인 T개 텐서)	(T*A, B) 또는 (A, T*B) (결합 차원에 따라)	(T, A, B) (dim=0일 때)
주요 사용 시나리오	동일한 속성의 데이터 배치 확장, RNN hidden state 연결, Transformer 결과 연결	독립적인 데이터들을 새로운 배치로 묶을 때, 여러 타임스텝의 데이터를 하나로 모을 때
비유	기차 칸 이어 붙이기	책상 위에 책 쌓기

---

4. 실무 해결 방법 7선: 개발자가 바로 적용 가능한 Example

실무에서 마주하는 구체적인 상황들과 그에 대한 torch.cat()과 torch.stack()을 활용한 해결 코드입니다. 이 예제들은 실무 수준의 복잡성을 반영하여 작성되었습니다.

Example 1: 동일한 형태의 이미지 배치 확장하기

가장 흔한 상황으로, 형태가 완벽하게 동일한 두 개의 이미지 배치를 하나로 합쳐 더 큰 배치를 만듭니다. 이때는 torch.cat()을 사용합니다.

import torch

# (Batch_Size, Channel, Height, Width) 형태의 텐서
batch1 = torch.randn(32, 3, 224, 224)
batch2 = torch.randn(16, 3, 224, 224)

# 기존 Batch 차원(dim=0)을 따라 연결
merged_batch = torch.cat([batch1, batch2], dim=0)

# 결과 형태: (48, 3, 224, 224)
print(f"Batch 1 shape: {batch1.shape}")
print(f"Batch 2 shape: {batch2.shape}")
print(f"Merged batch shape: {merged_batch.shape}")
# check rank: 4 == 4
print(f"Merged batch rank: {merged_batch.ndim}") 

Example 2: 독립적인 이미지들을 하나의 배치로 묶기

형태가 동일한 여러 장의 독립적인 이미지를 하나의 텐서로 묶어 모델의 입력으로 만듭니다. 새로운 배치 차원을 생성해야 하므로 torch.stack()을 사용합니다.

import torch

# 형태가 완벽하게 동일한 3장의 독립 이미지 (Channel, Height, Width)
img1 = torch.randn(3, 224, 224)
img2 = torch.randn(3, 224, 224)
img3 = torch.randn(3, 224, 224)

# 새로운 배치 차원(dim=0)을 생성하며 쌓음
image_batch = torch.stack([img1, img2, img3], dim=0)

# 결과 형태: (3, 3, 224, 224)
print(f"Single image shape: {img1.shape}")
print(f"Image batch shape: {image_batch.shape}")
# check rank: 3 -> 4
print(f"Image batch rank: {image_batch.ndim}") 

주의: 이미지들의 크기가 다르다면 torch.stack()은 불가능합니다. 전처리를 통해 크기를 맞추어야 합니다.

Example 3: RNN Hidden State 연결하기 (Bidirectional RNN)

Bidirectional RNN에서 정방향, 역방향 RNN의 hidden state를 연결(concatenate)하여 최종 hidden state를 만듭니다. 이때 각 hidden state는 (Batch, Hidden_Dim) 형태이며, hidden 차원(dim=1)을 따라 연결합니다.

import torch
import torch.nn as nn

batch_size = 16
hidden_dim = 64

# 정방향 hidden state
hidden_forward = torch.randn(batch_size, hidden_dim)
# 역방향 hidden state
hidden_backward = torch.randn(batch_size, hidden_dim)

# hidden 차원(dim=1)을 따라 연결
final_hidden = torch.cat([hidden_forward, hidden_backward], dim=1)

# 결과 형태: (16, 128)
print(f"Hidden forward shape: {hidden_forward.shape}")
print(f"Hidden backward shape: {hidden_backward.shape}")
print(f"Final hidden shape: {final_hidden.shape}")
# check rank: 2 == 2
print(f"Final hidden rank: {final_hidden.ndim}") 

Example 4: RNN 각 타임스텝의 Hidden State 모으기

RNN의 각 타임스텝에서 나온 hidden state들을 리스트에 모았다가, 하나의 텐서로 합칩니다. 각 hidden state는 (Batch, Hidden_Dim) 형태이며, 새로운 'Sequence' 차원을 생성하여 쌓아야 하므로 torch.stack()을 사용합니다.

import torch

batch_size = 16
hidden_dim = 64
seq_len = 10

# 각 타임스텝의 hidden state들을 담을 리스트
hidden_states = []

# 타임스텝 동안 hidden state 생성
for t in range(seq_len):
    # 각 타임스텝에서 생성된 hidden state: (Batch, Hidden_Dim)
    h_t = torch.randn(batch_size, hidden_dim)
    hidden_states.append(h_t)

# 새로운 Sequence 차원(dim=1)을 생성하며 쌓음
# 결과 형태: (Batch, Sequence, Hidden_Dim) -> (16, 10, 64)
merged_hidden = torch.stack(hidden_states, dim=1)

# 또는 dim=0으로 쌓아 (Sequence, Batch, Hidden_Dim) 형태로 만들 수도 있음 (PyTorch 표준)
merged_hidden_standard = torch.stack(hidden_states, dim=0)

print(f"Single hidden shape: {hidden_states[0].shape}")
print(f"Merged hidden (dim=1) shape: {merged_hidden.shape}")
print(f"Merged hidden standard (dim=0) shape: {merged_hidden_standard.shape}")
# check rank: 2 -> 3
print(f"Merged hidden rank: {merged_hidden.ndim}") 

Example 5: Transformer Multi-Head Attention 결과 연결하기

Transformer의 각 어텐션 헤드에서 나온 결과들을 어텐션 차원(dim=2)을 따라 연결합니다. 이때 각 헤드의 결과는 (Batch, Sequence, Head_Dim) 형태입니다.

import torch

batch_size = 16
seq_len = 50
head_dim = 64
num_heads = 8

# 각 어텐션 헤드의 결과 리스트
head_outputs = []

for _ in range(num_heads):
    # 각 헤드의 결과: (Batch, Sequence, Head_Dim)
    output_i = torch.randn(batch_size, seq_len, head_dim)
    head_outputs.append(output_i)

# 어텐션 차원(dim=2)을 따라 연결
multi_head_output = torch.cat(head_outputs, dim=2)

# 결과 형태: (16, 50, 512) (where 512 = 8 * 64)
print(f"Single head output shape: {head_outputs[0].shape}")
print(f"Multi-head output shape: {multi_head_output.shape}")
# check rank: 3 == 3
print(f"Multi-head output rank: {multi_head_output.ndim}") 

Example 6: Feature 차원의 크기가 다른 텐서 연결하기

Feature 차원의 크기가 다른 두 텐서를 Feature 차원(dim=1)을 따라 연결합니다. 이때 다른 차원의 크기는 동일해야 하므로 torch.cat()을 사용합니다.

import torch

batch_size = 16

# Feature 1: (Batch, Feature_Dim1)
feat1 = torch.randn(batch_size, 128)
# Feature 2: (Batch, Feature_Dim2)
feat2 = torch.randn(batch_size, 64)

# Feature 차원(dim=1)을 따라 연결
merged_feat = torch.cat([feat1, feat2], dim=1)

# 결과 형태: (16, 192)
print(f"Feature 1 shape: {feat1.shape}")
print(f"Feature 2 shape: {feat2.shape}")
print(f"Merged feature shape: {merged_feat.shape}")
# check rank: 2 == 2
print(f"Merged feature rank: {merged_feat.ndim}") 

Example 7: 텐서들을 합쳐 새로운 차원의 중간에 삽입하기

torch.stack()의 dim 인자를 활용해, 새로운 차원을 중간에 삽입할 수 있습니다. 예를 들어, (A, B) 형태의 텐서들을 합쳐 (A, T, B) 형태를 만듭니다.

import torch

a_dim = 5
b_dim = 10
t_count = 3

# (A, B) 형태의 텐서들
tensor1 = torch.randn(a_dim, b_dim)
tensor2 = torch.randn(a_dim, b_dim)
tensor3 = torch.randn(a_dim, b_dim)

# 새로운 차원을 dim=1 위치에 생성하며 쌓음
# 결과 형태: (5, 3, 10)
stacked_mid = torch.stack([tensor1, tensor2, tensor3], dim=1)

print(f"Original tensor shape: {tensor1.shape}")
print(f"Stacked mid shape: {stacked_mid.shape}")
# check rank: 2 -> 3
print(f"Stacked mid rank: {stacked_mid.ndim}") 

팁: dim=1은 새로운 차원이 결과 텐서의 index 1에 위치한다는 의미입니다. index 0의 차원(A)과 index 1의 차원(T)의 순서가 바뀐 것처럼 보이지만, 이는 torch.stack()의 정의에 따른 올바른 동작입니다.

---

5. 결론 및 전문가 조언

torch.cat()과 torch.stack()은 파이토치에서 텐서를 결합하는 두 가지 핵심적인 도구이지만, 그 동작 방식은 근본적으로 다릅니다.

• 기존 차원을 따라 데이터를 확장하고 싶다면 torch.cat()을 사용하십시오.
• 동일한 형태의 데이터들을 새로운 차원에 쌓아 묶고 싶다면 torch.stack()을 사용하십시오.

전문적인 관점에서, 모델 설계 시 차원의 변화를 명확히 이해하는 것은 에러 없는 코딩을 위한 첫걸음입니다. 위에서 제공한 비교표와 7가지 실무 예제를 완벽히 이해하고 활용한다면, 여러분은 차원Mismatch라는 지긋지긋한 문제로부터 자유로워질 수 있습니다. 텐서의 형태를 항상 확인하는 습관을 들이고, 이 글을 레퍼런스로 활용하여 더 견고하고 효율적인 파이토치 코드를 작성하시길 바랍니다.

내용 출처:
1. 파이토치(PyTorch) 공식 문서 - `torch.cat`, `torch.stack`.
2. 딥러닝 실무 경험 및 Transformer/RNN 모델 아키텍처 구현 사례.
3. 최신 인공지능 연구 논문 및 오픈 소스 코드 분석.