PyTorch中如何使LSTM模型输出的元素总和为1

你现在模型最后一层用的是Sigmoid()激活，这个函数只会把每个输出值单独压到0到1的区间，不会约束所有输出的总和，自然会出现你现在一堆0.5、总和远大于1的情况。

根据你的任务类型选对应的处理方式就行：

• 如果是单标签多分类任务（4433个类别里每个样本只属于一个类，需要输出每个类的概率，要求总和为1）：直接把最后一层的Sigmoid()换成Softmax()就可以。记得指定计算维度，一般这种结构最后一维是类别维，所以写torch.nn.Softmax(dim=-1)就行，它会沿着类别维度做归一化，保证每个样本的4433个输出加起来严格等于1，而且内部做了防溢出处理，数值稳定性有保障。
  改完的模型结构如下：

  LSTM(4433, 64)
  LSTM(64, 64)
  Linear(64, 4433)
  Softmax(dim=-1)
  

  踩坑提醒：如果你训练时用的是nn.CrossEntropyLoss当损失，这个损失内部已经自带Softmax计算了，这种情况模型结构里别加Softmax层，只在推理的时候对Linear层的原始输出调用torch.softmax()拿归一化后的概率就行，不然重复计算Softmax会导致模型效果下降。

• 要是你有特殊需求必须保留Sigmoid的输出逻辑，同时又要总和为1：可以在Sigmoid输出之后手动做归一化，代码写output = output / output.sum(dim=-1, keepdim=True)就行，本质是把每个输出值按占总输出的比例缩放，最后总和也能到1，但数值稳定性不如原生Softmax，没特殊需求不推荐使用。
• 如果你做的是多标签分类任务（每个样本可以同时属于多个类别）：那输出本来就不需要总和为1，现在用Sigmoid是正确选择，不用修改。

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