PyTorch带注意力RNN在均衡数据集上始终预测同一类别

模型塌缩到单类预测和数据集类别均衡性无关，完全是代码中的维度错误、激活函数搭配错误、逻辑错误导致的，具体问题和修复方案如下：

核心错误点 #

• 维度传递全链路错乱
  1. 注意力层传参完全错误：调用注意力层时传入的第二个参数是整数seq_length，但注意力层的第二个入参是mask，要求是形状为(batch_size, seq_len)的0/1掩码张量，传整数会导致注意力权重计算完全失效
  2. 注意力输出维度处理错误：注意力层的返回值是对整个序列加权求和后的单向量，形状为(batch_size, hidden_dim*2)，强行用view把它拆成三维序列张量喂给后续LSTM，会导致后续RNN拿到无意义的乱序特征
  3. 第二层GRU输入维度不匹配：取lstm_out2[:, -1]得到的是二维张量，而PyTorch的GRU默认要求输入是三维(batch, seq_len, feature)格式，直接传二维张量会导致输出形状完全失控
  4. embedding层多余维度操作：embeds = torch.squeeze(torch.unsqueeze(embeds, 0))这行没有任何实际作用，batch_size为1时还会错误挤压掉batch维度，引发维度错位
  5. 定义的Dropout层全程没有调用，完全没起到正则化作用
• 激活函数与损失函数搭配错误，梯度异常
  1. 注意力输出后接Sigmoid把特征压缩到0-1区间，再喂给RNN会导致梯度快速衰减，参数学不到有效信息
  2. 全连接层后接Softmax再喂给最终输出层：Softmax会把特征归一化为和为1的概率值，再经过线性变换后logits会被压缩到极窄区间，极易塌缩到固定类别
  3. 多分类任务用CrossEntropyLoss时，输出层前绝对不能加Softmax——CrossEntropyLoss内部已经集成了LogSoftmax计算，手动加Softmax会导致损失计算错误，反向传播梯度完全不对
• 注意力层缺乏数值稳定性处理
  直接对eij做指数运算torch.exp(eij)，当eij值稍大时会触发指数溢出，产生NaN/inf值，导致参数更新崩溃。

修复后代码 #

修正后的注意力层 #

class Attention(nn.Module):
    def __init__(self, feature_dim, bias=True, **kwargs):
        super(Attention, self).__init__(**kwargs)
        
        self.supports_masking = True
        self.bias = bias
        self.feature_dim = feature_dim
        
        weight = torch.zeros(feature_dim, 1)
        nn.init.kaiming_uniform_(weight)
        self.weight = nn.Parameter(weight)
        
        if bias:
            self.b = nn.Parameter(torch.zeros(1))
        
    def forward(self, x, mask=None):
        eij = torch.matmul(x, self.weight).squeeze(-1)
        
        if self.bias:
            eij = eij + self.b
            
        eij = torch.tanh(eij)
        # 数值稳定处理：减最大值再算指数，避免溢出
        eij = eij - torch.max(eij, dim=1, keepdim=True).values
        a = torch.exp(eij)
        
        if mask is not None:
            a = a * mask
        a = a / (torch.sum(a, 1, keepdim=True) + 1e-10)
        weighted_input = x * a.unsqueeze(-1)
        return torch.sum(weighted_input, 1)

修正后的分类模型 #

class AttentionLSTM(nn.Module):
    def __init__(self, weights_matrix, output_size, hidden_dim, drop_prob):
        super(AttentionLSTM, self).__init__()
        
        # embedding层
        self.embedding, self.num_embeddings, self.embeddings_size = create_emb_layer(weights_matrix, True)
        self.dropout = nn.Dropout2d(drop_prob)
        
        # 第一层RNN
        self.lstm1 = nn.LSTM(self.embeddings_size, hidden_dim, batch_first=True, bidirectional=True)
        self.gru1 = nn.GRU(hidden_dim * 2, hidden_dim, bidirectional=True, batch_first=True)
        
        # 注意力层无需提前固定seq_length
        self.attention = Attention(hidden_dim*2)

        # 注意力输出是单样本向量，无需再接序列RNN，直接接全连接层
        self.fc = nn.Linear(hidden_dim * 2, hidden_dim * 2)
        self.out = nn.Linear(hidden_dim * 2, output_size)
        self.relu = nn.ReLU()
        
    def forward(self, x, mask=None):
        x = x.long()
        embeds = self.embedding(x)
        embeds = self.dropout(embeds)
        
        lstm_out1, _ = self.lstm1(embeds)
        gru_out1, _ = self.gru1(lstm_out1)

        attention_out = self.attention(gru_out1, mask=mask)
        
        fc_out = self.relu(self.fc(attention_out))
        final_out = self.out(fc_out)
        # 直接返回原始logits，不要加Softmax
        return final_out

训练注意事项 #

• 损失函数直接用nn.CrossEntropyLoss()，不要对模型输出做Softmax后再计算损失
• 如果输入序列做过padding，提前生成对应mask（有效token位置为1，padding位置为0）传入forward，无padding则mask传None即可
• 初始学习率设为1e-3，用Adam优化器即可，先跑3-5个epoch观察训练损失是否稳定下降，如果仍不收敛，检查数据预处理阶段标签编码是否正确、输入张量是否存在大量NaN或全0值。

内容的提问来源于stack exchange，提问作者Büşra Şebin