### 如何在 PyTorch LSTM 中自定义激活函数 在 PyTorch 中，LSTM 的默认激活函数是 `tanh` 作为隐藏状态的激活函数和 `sigmoid` 作为门控机制的激活函数。如果需要自定义这些激活函数，可以通过继承 `nn.LSTM` 或直接修改源代码实现[^1]。 以下是一个通过自定义激活函数的方法示例： #### 方法一：重写 LSTM 的前向传播 可以继承 `nn.LSTM` 并覆盖其 `forward` 方法，从而替换默认的激活函数为自定义的激活函数。 ```python import torch import torch.nn as nn import torch.nn.functional as F class CustomLSTM(nn.LSTM): def __init__(self, input_size, hidden_size, num_layers=1, bias=True, batch_first=False, dropout=0, bidirectional=False, custom_activation=F.relu): super(CustomLSTM, self).__init__(input_size, hidden_size, num_layers, bias, batch_first, dropout, bidirectional) self.custom_activation = custom_activation # 自定义激活函数 def forward(self, input, hx=None): # 调用父类的 forward 方法获取输出 output, (h_n, c_n) = super().forward(input, hx) # 对输出应用自定义激活函数 output = self.custom_activation(output) return output, (h_n, c_n) ``` 在这个例子中，`custom_activation` 参数允许用户传入任何激活函数（例如 ReLU、Sigmoid 或其他自定义函数）。这样可以在每个时间步的输出上应用该激活函数。 #### 方法二：手动实现 LSTM 单元 如果需要更灵活地控制门控机制中的激活函数，可以手动实现 LSTM 单元，并完全自定义每个部分的激活函数。 ```python class CustomLSTMCell(nn.Module): def __init__(self, input_size, hidden_size, gate_activation=torch.sigmoid, cell_activation=torch.tanh): super(CustomLSTMCell, self).__init__() self.input_size = input_size self.hidden_size = hidden_size self.gate_activation = gate_activation # 门控激活函数 self.cell_activation = cell_activation # 细胞状态激活函数 # 定义权重矩阵 self.weight_ih = nn.Parameter(torch.randn(4 * hidden_size, input_size)) self.weight_hh = nn.Parameter(torch.randn(4 * hidden_size, hidden_size)) self.bias_ih = nn.Parameter(torch.randn(4 * hidden_size)) self.bias_hh = nn.Parameter(torch.randn(4 * hidden_size)) def forward(self, x, states): h_t, c_t = states gates = torch.mm(x, self.weight_ih.t()) + self.bias_ih + torch.mm(h_t, self.weight_hh.t()) + self.bias_hh i_t, f_t, g_t, o_t = gates.chunk(4, 1) # 应用门控激活函数 i_t = self.gate_activation(i_t) f_t = self.gate_activation(f_t) o_t = self.gate_activation(o_t) # 应用细胞状态激活函数 g_t = self.cell_activation(g_t) c_next = f_t * c_t + i_t * g_t h_next = o_t * self.cell_activation(c_next) return h_next, c_next ``` 在此方法中，`gate_activation` 和 `cell_activation` 是可配置的激活函数，分别用于门控机制和细胞状态更新。 #### 注意事项 - 在深层网络中，选择合适的激活函数非常重要。对于 LSTM，通常建议使用不会导致梯度消失的激活函数，如 ReLU 或 LeakyReLU[^3]。 - 如果选择非标准激活函数，需确保其导数特性适合序列建模任务，以避免训练过程中出现梯度爆炸或梯度消失问题[^2]。 ```python # 使用示例 lstm = CustomLSTM(input_size=10, hidden_size=20, custom_activation=torch.relu) x = torch.randn(5, 3, 10) # 输入形状为 (seq_len, batch, input_size) output, (h_n, c_n) = lstm(x) print(output.shape) # 输出形状为 (seq_len, batch, hidden_size) ```