[pytorch] LSTM 입력 텐서와 출력 텐서의 shape 이해

pytorch에서 LSTM 입출력 텐서의 shape 때문에 애를 먹었습니다. 입출력 텐서의 shape에 대해서 명확히 인지하고 있어야 모델 구성시 에러를 최소화 할 수 있습니다.

이번 튜토리얼 에서는 꽤나 복잡한 LSTM의 입출력 텐서 shape에 대해서 알아보도록 하겠습니다.

LSTM

• input_size: input x에 대한 features의 수

• hidden_size: hidden state의 features의 수

• num_layers: LSTM을 스택킹(stacking)하는 수

• batch_first: 입출력 텐서의 형태가 다음과 같음. 기본값은 False

  • True로 설정시 (batch, seq, feature)

  • False로 설정시 (seq, batch, feature)

  • (주의사항) hidden_state, cell_state에는 영향을 미치지 않습니다.

• bidirectional: 양방향 LSTM 구현 여부. 기본값은 False

import torch
import torch.nn as nn

lstm = nn.LSTM(input_size=10, hidden_size=20, num_layers=1, bidirectional=False, batch_first=True)
lstm

LSTM(10, 20, batch_first=True)

출력에 대한 이해

입력에 활용할 텐서 생성

inputs = torch.zeros(1, 35, 10)
print(inputs.shape)
# batch_size, sequence_length, number of features

torch.Size([1, 35, 10])

lstm에 방금 생성한 inputs를 입력으로 넣습니다.

여기서, lstm을 생성할 때 batch_first를 True로 설정했기 때문에 입력 텐서의 첫 번째 shape는 batch_size가 위치해야 합니다.

outputs, (hidden_state, cell_state) = lstm(inputs)
print(outputs.shape, hidden_state.shape, cell_state.shape)

torch.Size([1, 35, 20]) torch.Size([1, 1, 20]) torch.Size([1, 1, 20])

outputs 의 shape에 대한 이해

• [1, 35, 20] 형태를 가집니다. 순서대로 (batch_size, sequence_length, hidden_size(20) x bidirectional(1)) 입니다.

hidden_state 의 shape에 대한 이해

• hidden_state는 short term memory에 대한 state를 나타냅니다.

• [1, 1, 20] 형태를 가집니다. 순서대로 (Bidirectional x number of layers, batch_size, hidden_size) 입니다.

cell_state 의 shape에 대한 이해

• cell_state는 long term memory에 대한 state를 나타냅니다.

• [1, 1, 20] 형태를 가집니다. 순서대로 (Bidirectional x number of layers, batch_size, hidden_size) 입니다.

Stacking LSTM

• LSTM를 겹층으로 쌓아올리는 것을 Stacking이라고 합니다. 이렇게 쌓아올렸을 때는 결과 shape가 달라지게 됩니다.

• num_layers에 쌓고자 하는 층의 개수를 입력합니다.

# num_layers를 2로 설정
lstm = nn.LSTM(input_size=10, hidden_size=20, num_layers=2, bidirectional=False, batch_first=True)
lstm

LSTM(10, 20, num_layers=2, batch_first=True)

이전과 동일한 입력 텐서 생성

inputs = torch.zeros(1, 35, 10)
print(inputs.shape)
# batch_size, sequence_length, number of features

torch.Size([1, 35, 10])

outputs, (hidden_state, cell_state) = lstm(inputs)
print(outputs.shape, hidden_state.shape, cell_state.shape)

torch.Size([1, 35, 20]) torch.Size([2, 1, 20]) torch.Size([2, 1, 20])

outputs 의 shape에 대한 이해

• [1, 35, 20] 형태를 가집니다. 순서대로 (batch_size, sequence_length, hidden_size(20) x bidirectional(1)) 입니다.

hidden_state 의 shape에 대한 이해

• hidden_state는 short term memory에 대한 state를 나타냅니다.

• [2, 1, 20] 형태를 가집니다. 순서대로 (Bidirectional x number of layers, batch_size, hidden_size) 입니다.

• 출력 텐서의 첫 번째 shape은 Bidirectional(1) x number of layers(2)=2 결과입니다.

cell_state 의 shape에 대한 이해

• cell_state는 long term memory에 대한 state를 나타냅니다.

• [2, 1, 20] 형태를 가집니다. 순서대로 (Bidirectional x number of layers, batch_size, hidden_size) 입니다.

• 출력 텐서의 첫 번째 shape은 Bidirectional(1) x number of layers(2)=2 결과입니다.

Bidirectional LSTM

• 양방향 LSTM을 사용하기 위해서는 bidirectional=True로 설정합니다.

# bidirectional=True 설정
lstm = nn.LSTM(input_size=10, hidden_size=20, num_layers=2, bidirectional=True, batch_first=True)
lstm

LSTM(10, 20, num_layers=2, batch_first=True, bidirectional=True)

이전과 동일한 입력 텐서 생성

inputs = torch.zeros(1, 35, 10)
print(inputs.shape)
# batch_size, sequence_length, number of features

torch.Size([1, 35, 10])

outputs, (hidden_state, cell_state) = lstm(inputs)
print(outputs.shape, hidden_state.shape, cell_state.shape)

torch.Size([1, 35, 40]) torch.Size([4, 1, 20]) torch.Size([4, 1, 20])

outputs 의 shape에 대한 이해

• [1, 35, 40] 형태를 가집니다. 순서대로 (batch_size, sequence_length, hidden_size(20) x bidirectional(2)) 입니다.

• 출력 텐서의 마지막 shape은 hidden_size(20) x bidirectional(2)=40 결과입니다.

hidden_state 의 shape에 대한 이해

• hidden_state는 short term memory에 대한 state를 나타냅니다.

• [4, 1, 20] 형태를 가집니다. 순서대로 (bidirectional x number of layers, batch_size, hidden_size) 입니다.

• 출력 텐서의 첫 번째 shape은 bidirectional(2) x number of layers(2)=4 결과입니다.

cell_state 의 shape에 대한 이해

• cell_state는 long term memory에 대한 state를 나타냅니다.

• [4, 1, 20] 형태를 가집니다. 순서대로 (bidirectional x number of layers, batch_size, hidden_size) 입니다.

• 출력 텐서의 첫 번째 shape은 bidirectional(2) x number of layers(2)=4 결과입니다.