[PyTorch] nn.LSTM input, output shape

LSTM

import torch
from torch import nn

input = torch.randn(5, 3, 10) # (seq_len, batch, input_size)

 input 데이터는 3개, LSTM에 순차적으로 들어갈 시간은 5개, 입력 노드의 개수(설명변수)는 10개라고 하자. torch.randn로 rnn에 들어갈 input을 만들면 첫 번째 원소가 sequence length = 5, 두 번째 원소가 batch size = 3, 세 번째 원소가 input size = 10 인 차원으로 생성한다.

 PyTorch에서 다음과 같이 sequence 길이는 5, 출력 노드의 개수는 20개, LSTM층이 2개인 RNN 구조를 만들어보자.

 

 nn.LSTM에서 입력층의 크기(input_size)는 10 (주황색 화살표), 은닉층의 크기(hidden_size)는 20 (파란색 화살표), 레이어의 개수(num_layers)는 2개인 순환신경망이다. input 데이터를 rnn에 집어 넣으면 output과 (hn, cn)의 튜플을 반환한다.

rnn = nn.LSTM(input_size = 10, hidden_size = 20, num_layers = 2)
output, (hn, cn) = rnn(input)

 

output은 2층 layer의 시간 순서에 따른 출력이다. 첫 번째 원소는 sequence 길이인 5이고, 두 번째 원소는 batch 크기, 세 번째 원소는 20으로 은닉층의 크기(출력층의 크기)이다.

output.size()

> torch.Size([5, 3, 20])

 

 

 

다음과 같이 last sequence의 출력만 사용하려면 output에서 마지막 값을 가져온다. output[-1]의 shape를 확인해보면 batch가 3개, 은닉층(출력)의 크기가 20이 됨을 알 수 있다.

output[-1].size()

> torch.Size([3, 20])

 

 

 다음으로 (hn, cn)은 각각 다음 LSTMcell로 넘겨주는 hidden state, cell sate를 말하고, input으로 (hn, cn)이 주어지지 않으면 c0와 h0는 zeros로 초기화한다. hn과 cn의 shape를 살펴보자.

hn과 cn의 차원의 첫 번째 원소는 2층 레이어이므로 2, 두 번째원 소는 batch 3개, 은닉층의 크기이므로 20이다.

print(hn.size())
print(cn.size())

> torch.Size([2, 3, 20])

> torch.Size([2, 3, 20])

 

 

Bidirectional LSTM

nn.LSTM에서는 bidirectional = True로 설정하므로써 간단하게 위와 같은 biderectional LSTM을 구현할 수 있다.

import torch
from torch import nn

input = torch.randn(5, 3, 10) # (seq_len, batch, input_size)
rnn = nn.LSTM(input_size = 10, hidden_size = 20, num_layers = 2, bidirectional = True)
output, (hn, cn) = rnn(input)

 이 때 output은 왼쪽에서 오른쪽으로 한 개, 오른쪽에서 왼쪽으로 한 개가 concatenate 되어 출력이 된다. 따라서 output의 hidden_size(출력층의 크기)는 두 배인 2 * 20 = 40이 된다.

 

output.size()

> torch.Size([5, 3, 40])

 

 

 hn과 cn의 첫 번째 원소는 bidirection이므로 2개 * 2층 레이어이므로 2개 = 총 4개이다. 다음 LSTM cell로 넘겨주는 hn과 cn은 hidden_size가 concatenate 되지 않으므로 20인 것을 알 수 있다.

print(hn.size())
print(cn.size())

> torch.Size([4, 3, 20])

> torch.Size([4, 3, 20])