如何在PyTorch中构建输出固定时序参数的LSTM混合水文模型

我在PyTorch中构建了融合概念性水文模型与LSTM的混合模型：LSTM输出参数，水文模型结合这些参数与降水数据生成时序河流流量。目前需要实现**输出固定参数（参数不随时间步变化）**的LSTM，但当前实现方式较为笨拙，附上模型代码片段，求更合理的实现方案。

原模型代码 #

# 参数配置
#network_params = {
#  "input_size": len(attributes) + len(forcing) - 1,  # 输入维度（属性数+驱动变量数-日期）
#  "no_of_layers": 1,  # LSTM层数
#  "sequence_length": 180,  # LSTM序列长度
#  "warmup_period": 365,  # 预热周期
#  "batches_to_update": 20,  # 更新参数的批次数量
#  "hidden_size": 64,  # 隐藏层维度
#  "no_of_epochs": 20,  # 训练轮数
#  "drop_out": 0.4,  # Dropout比例
#  "learning_rate": 0.01,  # 学习率
#  'set_forget_gate': 3,  # 遗忘门配置
#  "adapt_learning_rate_epoch": 5,  # 调整学习率的轮数
#  "adapt_gamma_learning_rate": 0.8  # 学习率调整因子
#}

# 混合模型类：融合LSTM与水文模型
class Hybrid_Model(nn.Module):
    
    def __init__(self, network_params, cuda):
        
        super().__init__()

        self.num_features = network_params['input_size']
        self.hidden_units = network_params['hidden_size']
        self.num_layers   = network_params['no_of_layers']
        self.cuda0 = cuda

        # 构建LSTM网络
        self.lstm = nn.LSTM(input_size = network_params['input_size'], 
                            hidden_size = network_params['hidden_size'], 
                            batch_first = True,
                            num_layers = network_params['no_of_layers'])

        self.dropout = torch.nn.Dropout(network_params['drop_out'])
        
        # 输出层维度需根据参数数量调整
        self.linear = nn.Linear(in_features=network_params['hidden_size'], out_features=4)
        
        # 构建水文模型
        self.SHM = SHM_dynamic_parameters()
           
    def forward(self, X_LSTM, X_SHM, initial_states, warmup_period=0):
        
        # 初始化隐藏状态
        batch_size = X_LSTM.shape[0]
        
        h0 = torch.zeros(self.num_layers, batch_size, self.hidden_units).requires_grad_().to(self.cuda0)
        c0 = torch.zeros(self.num_layers, batch_size, self.hidden_units).requires_grad_().to(self.cuda0)
        
        # LSTM运行逻辑：一次性预测水文模型所需的所有时间步参数
        _, (hn, _) = self.lstm(X_LSTM, (h0, c0))
        
        out = hn[-1]
        
        out = self.dropout(out)
        out = self.linear(out)
        out = torch.sigmoid(out)
        
        out_copy = out.clone()

        # 预热阶段：稳定水文模型内部状态（储水桶）
        if warmup_period > 0:
            with torch.no_grad():
                _, _, initial_states = self.SHM(X_SHM=X_SHM[0:warmup_period, :], 
                                                 out_LSTM=out_copy.unsqueeze(1),  # 添加时间维度
                                                 initial_states=initial_states, 
                                                 cuda=self.cuda0)
                
                initial_states = initial_states[-1, :]
                
        # 运行模型
        q_out, parameters, new_states = self.SHM(X_SHM=X_SHM[warmup_period:, :], 
                                                 out_LSTM=out_copy.unsqueeze(1),  # 添加时间维度
                                                 initial_states=initial_states, 
                                                 cuda=self.cuda0)

        return q_out, parameters, new_states

1. 简化固定参数生成逻辑 #

当前通过hn[-1]获取LSTM最后一层的最终隐藏状态，本质已经是不随时间步变化的全局特征，这部分逻辑合理，但可以优化参数维度扩展的方式：直接计算水文模型需要的时间步数，用repeat明确扩展时间维度，替代unsqueeze(1)的隐式操作，可读性更强。

2. 移除不必要的张量克隆 #

原代码中out_copy = out.clone()属于冗余操作，后续仅对张量做维度扩展，无原地修改，直接使用out即可。

3. 优化隐藏状态初始化 #

将隐藏状态的零张量初始化移至__init__方法，用register_buffer创建可复用的缓冲区，避免每次forward都重复创建张量，提升效率。

4. 可选：用MLP替代LSTM（如果序列仅用于提取全局特征） #

如果输入序列的作用只是提取全局统计特征而非时序依赖，可直接用全局池化+线性层替代LSTM，计算更高效。

class Hybrid_Model(nn.Module):
    
    def __init__(self, network_params, cuda):
        
        super().__init__()

        self.num_features = network_params['input_size']
        self.hidden_units = network_params['hidden_size']
        self.num_layers   = network_params['no_of_layers']
        self.cuda0 = cuda

        # 构建LSTM网络
        self.lstm = nn.LSTM(input_size = network_params['input_size'], 
                            hidden_size = network_params['hidden_size'], 
                            batch_first = True,
                            num_layers = network_params['no_of_layers'])

        self.dropout = torch.nn.Dropout(network_params['drop_out'])
        self.linear = nn.Linear(in_features=network_params['hidden_size'], out_features=4)
        
        # 注册隐藏状态初始化缓冲区，避免重复创建
        self.register_buffer('h0_init', torch.zeros(self.num_layers, 1, self.hidden_units))
        self.register_buffer('c0_init', torch.zeros(self.num_layers, 1, self.hidden_units))
        
        # 构建水文模型
        self.SHM = SHM_dynamic_parameters()
           
    def forward(self, X_LSTM, X_SHM, initial_states, warmup_period=0):
        
        batch_size = X_LSTM.shape[0]
        # 扩展初始化隐藏状态到当前批次大小
        h0 = self.h0_init.expand(self.num_layers, batch_size, self.hidden_units).requires_grad_().to(self.cuda0)
        c0 = self.c0_init.expand(self.num_layers, batch_size, self.hidden_units).requires_grad_().to(self.cuda0)
        
        # LSTM输出固定参数
        _, (hn, _) = self.lstm(X_LSTM, (h0, c0))
        out = hn[-1]
        out = self.dropout(out)
        out = self.linear(out)
        out = torch.sigmoid(out)

        # 预热阶段
        if warmup_period > 0:
            warmup_steps = X_SHM[:warmup_period].shape[0]
            # 明确扩展参数到预热时间步维度
            warmup_params = out.unsqueeze(1).repeat(1, warmup_steps, 1)
            with torch.no_grad():
                _, _, initial_states = self.SHM(X_SHM=X_SHM[:warmup_period], 
                                                 out_LSTM=warmup_params,
                                                 initial_states=initial_states, 
                                                 cuda=self.cuda0)
                initial_states = initial_states[-1]
                
        # 正式运行阶段
        run_steps = X_SHM[warmup_period:].shape[0]
        run_params = out.unsqueeze(1).repeat(1, run_steps, 1)
        q_out, parameters, new_states = self.SHM(X_SHM=X_SHM[warmup_period:], 
                                                 out_LSTM=run_params,
                                                 initial_states=initial_states, 
                                                 cuda=self.cuda0)

        return q_out, parameters, new_states

内容的提问来源于stack exchange，提问作者Vinicius B. de S. Moreira