PyTorch nn.LSTM 各层节点数按系数递减的实现方法咨询

原生torch.nn.LSTM确实没有提供逐层配置不同隐藏维度的入参，其num_layers参数对应的所有堆叠层会统一使用传入的hidden_size值，因此天然不支持等比/固定系数递减的非对称维度堆叠。
要实现无需手动逐一定义LSTM层的效果，最简洁的方式是通过动态维度计算+nn.ModuleList批量生成层实例的方案，全程不需要手写每一层的LSTM初始化代码，生成的层结构能和原有代码无缝兼容。

具体实现逻辑 #

• 维度动态计算：基于初始隐藏维度、固定递减系数、总层数，自动推导每一层的输入、输出维度，无需手动罗列
  • 第0层的输入维度为模型全局的input_size，输出维度为初始隐藏维度
  • 第i层（i≥1）的输入维度为上一层的输出维度，输出维度按固定系数做递减计算
• 层实例动态生成：通过循环把每一层的LSTM存入nn.ModuleList，确保PyTorch能正确追踪所有可训练参数
• 层间逻辑对齐：在相邻LSTM层之间按原生LSTM的逻辑插入Dropout层，保证和原生实现的训练行为完全一致

可直接复用的代码实现 #

替换原有LSTM初始化逻辑的完整类代码如下，调用方式和原生nn.LSTM几乎没有区别：

import torch.nn as nn
import math

class DecayLSTM(nn.Module):
    def __init__(self, input_size, init_hidden_size, num_layers, decay_ratio, dropout=0.2, batch_first=True, bidirectional=False):
        super().__init__()
        # 用ModuleList存储所有动态生成的LSTM层
        self.lstm_layers = nn.ModuleList()
        self.bidirectional = bidirectional
        current_input_dim = input_size
        
        for layer_idx in range(num_layers):
            # 按固定系数计算当前层的隐藏维度，向下取整保证为整数
            current_hidden_dim = math.floor(init_hidden_size * (decay_ratio ** layer_idx))
            self.lstm_layers.append(
                nn.LSTM(
                    input_size=current_input_dim,
                    hidden_size=current_hidden_dim,
                    num_layers=1,
                    batch_first=batch_first,
                    bidirectional=bidirectional,
                    dropout=0  # 关闭单层内置dropout，统一在层间添加
                )
            )
            # 更新下一层的输入维度：如果是双向LSTM，输出维度为当前隐藏维度的2倍
            current_input_dim = current_hidden_dim * (2 if bidirectional else 1)
        
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        h_n_final, c_n_final = None, None
        for idx, lstm in enumerate(self.lstm_layers):
            x, (h_n, c_n) = lstm(x)
            # 非最后一层的输出后添加dropout，和原生LSTM行为对齐
            if idx != len(self.lstm_layers) - 1:
                x = self.dropout(x)
            # 保留最后一层的隐状态作为最终输出
            h_n_final, c_n_final = h_n, c_n
        return x, (h_n_final, c_n_final)

调用方式 #

完全不需要修改原有前向传播逻辑，仅需替换原来的LSTM初始化代码即可：

# 原代码：self.lstm = nn.LSTM(...)
# 替换为下面的初始化，参数含义直观：
self.lstm = DecayLSTM(
    input_size=input_size,
    init_hidden_size=hidden_size,  # 首层的隐藏维度
    num_layers=num_layers,
    decay_ratio=0.5,  # 每层隐藏维度相比上一层的递减系数
    dropout=0.2,
    batch_first=True
)

关键注意事项 #

• 禁止用普通Python列表存储LSTM实例：如果使用普通list存储生成的层，PyTorch的参数追踪机制会无法识别这些层的可训练权重，导致训练时参数不更新，必须使用nn.ModuleList做存储
• 维度取整处理：按固定系数计算出的隐藏维度可能为小数，用math.floor向下取整即可，不会对模型效果产生显著影响
• 输出形状完全对齐：该实现的输出张量形状、隐状态形状和原生nn.LSTM保持一致，原有前向传播的后续逻辑不需要做任何适配修改

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