torch.cat vs torch.stack

torch.cat 과 torch.stack 은 서로 다른 텐서 결합 방법이다. 

 

1) torch.cat (Concatenate)

torch.cat 은 서로 다른 행렬이었던 두 텐서를 하나의 텐서로 만들어준다. 따라서 차원의 변화가 생기지 않는다. 

예를 들어 두 개의 2x3 짜리 행렬을 cat 해주면 쌓아주는 방향에 따라 4x3 행렬이나 2x6 행렬이 된다. 

 

import torch

# 2x3 행렬 하나
a = torch.tensor([[1,2,3],[4,5,6]])
# 2x3 행렬 하나
b = torch.tensor([[7,8,9],[10,11,12]])

# 행방향 쌓기
c = torch.cat([a,b], dim=0)

# 열방향 쌓기
d = torch.cat([a,b], dim=1)

print(c)
print(d)

# c 텐서, 4x3 행렬 하나
tensor([[ 1,  2,  3],
        [ 4,  5,  6],
        [ 7,  8,  9],
        [10, 11, 12]])

# d 텐서, 2x6 행렬 하나
tensor([[ 1,  2,  3,  7,  8,  9],
        [ 4,  5,  6, 10, 11, 12]])

이렇듯 두 개의 tensor 를 cat하면 차원 변화 없이 하나의 텐서로 합쳐지는 것을 볼 수 있다. 

 

2) torch.stack

torch.stack 함수는 서로 다른 텐서들을 새로운 차원을 생성하여 연결해준다.

따라서 결과 텐서는 입력 텐서보다 차원이 하나 더 많은 경우가 많다. 예를 들어 2x3 행렬 두개를 torch.stack 하면 (2,2,3) 텐서가 된다.

torch.stack 은 행방향과 열방향으로 쌓았을 때 결과가 모두 (2,2,3) 이었는데,

이처럼 기존의 행렬 사이즈(2,3) 는 변화하지 않고, 차원만 하나 더 늘어나게 된다.

 

import torch

# 2x3 행렬 하나
a = torch.tensor([[1,2,3],[4,5,6]])
# 2x3 행렬 하나
b = torch.tensor([[7,8,9],[10,11,12]])

# 행방향 쌓기
c = torch.stack([a,b], dim=0)
# 열방향 쌓기
d = torch.stack([a,b], dim=1)

print(c)
print(d)

# c 텐서, 2x3 행렬 2개
tensor([[[ 1,  2,  3],
         [ 4,  5,  6]],

        [[ 7,  8,  9],
         [10, 11, 12]]])
         
# d 텐서, 2x3 행렬 2개
tensor([[[ 1,  2,  3],
         [ 7,  8,  9]],

        [[ 4,  5,  6],
         [10, 11, 12]]])

주의할점

torch.stack 이나 torch.cat 을 사용할 때는 함수의 인자로 리스트나 튜플을 전달해야 한다.

torch.stack(a,b)와 같이 그냥 텐서로만 넣어준다면 다음과 같은 오류가 발생할 것이다.

TypeError: stack(): argument 'tensors' (position 1) must be tuple of Tensors, not Tensor

 

따라서 꼭 튜플이나 리스트로 묶은 다음 알맞은 함수를 사용해야 한다.