pytorch 中LSTM模型獲取最后一層的輸出結果，單向或雙向

單向LSTM

 

import torch.nn as nn
import torch

seq_len = 20
batch_size = 64
embedding_dim = 100
num_embeddings = 300
hidden_size = 128
number_layer = 3

input = torch.randint(low=0,high=256,size=[batch_size,seq_len])  #[64,20]

embedding = nn.Embedding(num_embeddings,embedding_dim)

input_embeded = embedding(input)  #[64,20,100]

#轉置，變換batch_size 和seq_len
# input_embeded = input_embeded.transpose(0,1)
# input_embeded = input_embeded.permute(1,0,2)
#實例化lstm

lstm = nn.LSTM(input_size=embedding_dim,hidden_size=hidden_size,batch_first=True,num_layers=number_layer)

output,(h_n,c_n) = lstm(input_embeded)
print(output.size()) #[64,20,128]       [batch_size,seq_len,hidden_size]
print(h_n.size()) #[3,64,128]           [number_layer,batch_size,hidden_size]
print(c_n.size()) #同上


#獲取最后時間步的output
output_last = output[:,-1,:]
#獲取最后一層的h_n
h_n_last = h_n[-1]

print(output_last.size())
print(h_n_last.size())
#最后的output等於最后一層的h_n
print(output_last.eq(h_n_last))

D:\anaconda\python.exe C:/Users/liuxinyu/Desktop/pytorch_test/day4/LSTM練習.py
torch.Size([64, 20, 128])
torch.Size([3, 64, 128])
torch.Size([3, 64, 128])
torch.Size([64, 128])
torch.Size([64, 128])
tensor([[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
...,
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True]])

Process finished with exit code 0

 

　　雙向LSTM

 

import torch.nn as nn
import torch

seq_len = 20
batch_size = 64
embedding_dim = 100
num_embeddings = 300
hidden_size = 128
number_layer = 3

input = torch.randint(low=0,high=256,size=[batch_size,seq_len])  #[64,20]

embedding = nn.Embedding(num_embeddings,embedding_dim)

input_embeded = embedding(input)  #[64,20,100]

#轉置，變換batch_size 和seq_len
# input_embeded = input_embeded.transpose(0,1)
# input_embeded = input_embeded.permute(1,0,2)
#實例化lstm

lstm = nn.LSTM(input_size=embedding_dim,hidden_size=hidden_size,batch_first=True,num_layers=number_layer,bidirectional=True)

output,(h_n,c_n) = lstm(input_embeded)
print(output.size()) #[64,20,128*2]       [batch_size,seq_len,hidden_size]
print(h_n.size()) #[3*2,64,128]           [number_layer,batch_size,hidden_size]
print(c_n.size()) #同上


#獲取反向的最后一個output
output_last = output[:,0,-128:]
#獲反向最后一層的h_n
h_n_last = h_n[-1]

print(output_last.size())
print(h_n_last.size())
# 反向最后的output等於最后一層的h_n
print(output_last.eq(h_n_last))

#獲取正向的最后一個output
output_last = output[:,-1,:128]
#獲取正向最后一層的h_n
h_n_last = h_n[-2]
# 反向最后的output等於最后一層的h_n
print(output_last.eq(h_n_last))

D:\anaconda\python.exe C:/Users/liuxinyu/Desktop/pytorch_test/day4/雙向LSTM練習.py
torch.Size([64, 20, 256])
torch.Size([6, 64, 128])
torch.Size([6, 64, 128])
torch.Size([64, 128])
torch.Size([64, 128])
tensor([[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
...,
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True]])
tensor([[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
...,
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True],
[True, True, True, ..., True, True, True]])

Process finished with exit code 0