PyTorch中使用PackedSequence时如何处理LSTM层之间的Dropout？

我正在为硕士论文搭建用于特征提取的LSTM自编码器，目前在将Dropout与LSTM层结合使用时遇到了较多问题。

由于是自编码器结构，我需要设置瓶颈层，具体通过两个独立的、num_layers均为1的LSTM层，以及二者之间的Dropout层实现。我处理的时序数据长度差异极大，因此采用PackedSequence来适配这一特性。

但在实验中我发现，必须在输入第一个LSTM前打包数据，Dropout前解包，第二个LSTM前再次打包，这种方式效率很低，请问是否有更优的实现方案？下方是我目前的可运行但可能非最优的代码，以及一套尚未跑通的备选实现方案。

当前可运行但可能非最优的解决方案 #

class Encoder(nn.Module):

    def __init__(self, seq_len, n_features, embedding_dim, hidden_dim, dropout):
        super(Encoder, self).__init__()

        self.seq_len = seq_len
        self.n_features = n_features
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim

        self.lstm1 = nn.LSTM(
            input_size=n_features,
            hidden_size=self.hidden_dim,
            num_layers=1,
            batch_first=True,
        )

        self.lstm2 = nn.LSTM(
            input_size=self.hidden_dim,
            hidden_size=embedding_dim,
            num_layers=1,
            batch_first=True,
        )

        self.drop1 = nn.Dropout(p=dropout, inplace=False)

    def forward(self, x):
        x, (_, _) = self.lstm1(x)
        x, lens = pad_packed_sequence(x, batch_first=True, total_length=self.seq_len)
        x = self.drop1(x)
        x = pack_padded_sequence(x, lens, batch_first=True, enforce_sorted=False)
        x, (hidden_n, _) = self.lstm2(x)

        return hidden_n.reshape((-1, self.n_features, self.embedding_dim)), lens

备选的潜在更优但当前无法运行的解决方案 #

class Encoder2(nn.Module):

    def __init__(self, seq_len, n_features, embedding_dim, hidden_dim, dropout):
        super(Encoder2, self).__init__()

        self.seq_len = seq_len
        self.n_features = n_features
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim

        self.lstm1 = nn.LSTM(
            input_size=n_features,
            hidden_size=self.hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=dropout,
            proj_size=self.embedding_dim,
        )

    def forward(self, x):
        _, (h_n, _) = self.lstm1(x)
        return h_n[-1].unsqueeze(1), lens

如果有关于时序数据处理、PackedSequence、LSTM单元以及Dropout相关的建议也非常欢迎，我在网上几乎找不到相关的文档与指导，非常感谢大家的帮助！

此致
Lars Ankile

现有方案的直接优化 #

PyTorch 1.1.0及以上版本的nn.Dropout原生支持PackedSequence类型输入，完全不需要手动解包、应用Dropout再重新打包，你现有方案里的解包打包步骤可以直接删除，优化后的forward逻辑如下：

def forward(self, x):
    # x为输入的PackedSequence
    x, (_, _) = self.lstm1(x)
    # 直接对PackedSequence应用Dropout，无需解包
    x = self.drop1(x)
    x, (hidden_n, _) = self.lstm2(x)
    # 若需要返回序列长度，可直接从输入侧传入，无需解包获取
    return hidden_n.reshape((-1, self.n_features, self.embedding_dim)), x.batch_sizes

该优化不需要改动其他逻辑，即可大幅减少不必要的张量操作，提升运行效率。

备选方案的问题修复 #

你编写的Encoder2存在3个可修复的问题，调整后可正常运行，且效率比第一个方案更高：

1. 变量lens未定义：可直接将打包输入时的序列长度列表作为参数传入forward，无需从运算结果中提取
2. proj_size参数约束：当设置proj_size>0时，要求hidden_dim >= embedding_dim，否则会触发参数不合法报错
3. 隐藏态维度适配：开启proj_size后，多层LSTM的输出隐藏态维度为proj_size，适配你的瓶颈层输出要求

修复后的代码如下：

class Encoder2(nn.Module):
    def __init__(self, seq_len, n_features, embedding_dim, hidden_dim, dropout):
        super(Encoder2, self).__init__()
        self.seq_len = seq_len
        self.n_features = n_features
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        # 校验proj_size参数合法性
        assert hidden_dim >= embedding_dim, "使用proj_size时hidden_dim必须大于等于embedding_dim"
        self.lstm1 = nn.LSTM(
            input_size=n_features,
            hidden_size=self.hidden_dim,
            num_layers=2,
            batch_first=True,
            dropout=dropout,
            proj_size=self.embedding_dim,
        )

    def forward(self, x, lens):
        # x为输入的PackedSequence，lens为输入时的序列长度列表
        _, (h_n, _) = self.lstm1(x)
        return h_n[-1].unsqueeze(1), lens

该方案全程不需要任何解包打包操作，官方内部实现的层间Dropout也比手动添加的层效率更高，更适合处理大批次的可变长度时序数据。

额外建议 #

• 处理可变长度时序数据时，只要输入是合法的PackedSequence，PyTorch的所有RNN类层、Dropout、普通激活层都可以直接处理，只有当你需要对补零后的完整张量做维度操作时，才需要调用pad_packed_sequence
• 不要手动在LSTM的循环内部加dropout，官方提供的dropout参数是标准的层间dropout，完全符合大部分任务的正则需求；如果需要循环层面的变分dropout，可使用torch.nn的相关扩展实现
• 自编码器瓶颈层取最后一个隐藏态作为编码时，PackedSequence模式下LSTM返回的隐藏态已经是对应序列实际长度的最后一步结果，不会包含补零部分的干扰，不需要额外做掩码处理

内容的提问来源于stack exchange，提问作者larsankile