pytorch nn.LSTM()参数详解

函数

class torch.nn.LSTM(*args, **kwargs)

参数列表

• input_size：x的特征维度
• hidden_size：隐藏层的特征维度
• num_layers：lstm隐层的层数，默认为1
• bias：False则bih=0和bhh=0. 默认为True
• batch_first：True则输入输出的数据格式为 (batch, seq, feature)
• dropout：除最后一层，每一层的输出都进行dropout，默认为: 0
• bidirectional：True则为双向lstm默认为False
  输入：input, (h0, c0)
  输出：output, (hn,cn)

输入数据格式：

input(seq_len, batch, input_size)
h0(num_layers * num_directions, batch, hidden_size)
c0(num_layers * num_directions, batch, hidden_size)

输出数据格式：

output(seq_len, batch, hidden_size * num_directions)
hn(num_layers * num_directions, batch, hidden_size)
cn(num_layers * num_directions, batch, hidden_size)

num_directions指的是单向（值为1）还是双向（值为2）。

import torch
import torch.nn as nn
from torch.autograd import Variable

# 构建网络模型---输入矩阵特征数input_size、输出矩阵特征数hidden_size、层数num_layers
inputs = torch.randn(5,3,10)   # (seq_len,batch_size,input_size)
rnn = nn.LSTM(10,20,2)         # (input_size,hidden_size,num_layers)
h0 = torch.randn(2,3,20)       # (num_layers* 1,batch_size,hidden_size)
c0 = torch.randn(2,3,20)       # (num_layers*1,batch_size,hidden_size) 
num_directions=1               # 因为是单向LSTM

output,(hn,cn) = rnn(inputs,(h0,c0))

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

'''输出
torch.Size([5, 3, 20])
torch.Size([2, 3, 20])
torch.Size([2, 3, 20])
'''

batch_first 输入输出的第一维是否为 batch_size，默认值 False。因为 Torch 中，人们习惯使用Torch中带有的dataset，dataloader向神经网络模型连续输入数据，这里面就有一个 batch_size 的参数，表示一次输入多少个数据。 在 LSTM 模型中，输入数据必须是一批数据，为了区分LSTM中的批量数据和dataloader中的批量数据是否相同意义，LSTM 模型就通过这个参数的设定来区分。 如果是相同意义的，就设置为True，如果不同意义的，设置为False。 torch.LSTM 中 batch_size 维度默认是放在第二维度，故此参数设置可以将 batch_size 放在第一维度。如：input 默认是(4,1,5)，中间的 1 是 batch_size，指定batch_first=True后就是(1,4,5)。所以，如果你的输入数据是二维数据的话，就应该将 batch_first 设置为True;