pack_padded_sequence, pad_packed_sequence 설명

pack.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
# Embedding dim: E = 12
# Batch size: B = 2
# Tokens: T = 6
# Hidden Size: H = 20
# Vocab Size: V = 10
E, B, T, H, V = (12, 1, 6, 20, 10)
embedding_layer = nn.Embedding(V, E)
lstm = nn.LSTM(E, H, 1, batch_first=True, bidirectional=False)

# 1. 데이터 불러오기
# x: data loader 에서 나온 패딩이 된 데이터라고 하면
# 아래와 같은 코드가 있어야할 것입니다.
# for x in Dataloader:
#   ...
# 현재 배치 = 2, 패딩된 sequence 총 길이 = 6
# 즉, 배치의 1번 데이터는 길이가 6, 2번 데이터는 길이가 3임으로 
# 제일 긴 1번 데이터에 맞춰서 2번 데이터를 3개의 토큰을 패딩합니다.
x = torch.LongTensor([[3, 9, 2, 1, 5, 6],
                      [1, 2, 4, 0, 0, 0]])
# x: B, T

# 2. 모델 내부: Foward 과정
# output = model(x)

embed = embedding_layer(x)
# embed: B, T, E

# 패딩된 문장을 패킹(패딩은 연산 안들어가도록)
packed = pack_padded_sequence(embed, inputs_lengths.tolist(), batch_first=True)
# packed: B * T, E
init_hidden = torch.zeros(B, T, H)
init_cell = torch.zeros(B, T, H)
output, (hidden, cell) = lstm(packed, (hidden, cell))
# output: B * T, H
# hidden, cell: num_layers * num_directions, B, u

# 패킹된 문장을 다시 unpack
output, output_lengths = pad_packed_sequence(output, batch_first=True)
# output: B, T, H