PyTorch训练二分类LSTM出现Target与input张量尺寸不匹配报错

报错根因 #

你当前使用的BCEWithLogitsLoss是单输出二分类专用损失，要求模型输出和标签的形状完全一致，而你设置模型output_dim=2时输出形状为[1, 2]，标签经过unsqueeze(1)后形状为[1, 1]，尺寸不匹配触发报错。

解决方案二选一 #

方案1：保留output_dim=2，切换适配的损失函数 #

适合二分类用双输出逻辑的场景：

1. 把损失函数替换为交叉熵损失：

loss_func = nn.CrossEntropyLoss()

2. 删除训练代码中多余的标签维度调整代码，删掉两行：

# 删掉这两行
Y_tr = Y_tr.type(torch.LongTensor)
Y_tr = Y_tr.unsqueeze(1)

3. 修改损失计算逻辑：

loss = loss_func(output, Y_tr.long())

方案2：设置output_dim=1，继续使用BCEWithLogitsLoss #

就是你之前尝试的方向，修改对应问题即可解决恒预测同一类的问题：

1. 首先修改模型输出维度为1，此时模型输出形状为[1,1]，和标签形状匹配，不会再报尺寸错误
2. 修复恒预测同一类的问题，你的错误出在三个地方：

• 预测逻辑错误：单输出时不能用argmax取结果，[1,1]形状的张量argmax(dim=1)永远返回索引0，所以看起来永远预测同一类，替换为：

# 单输出logits大于0即为正类
pred = (output > 0).long()

• LSTM配置错误：nn.LSTM的dropout参数仅在层数≥2时生效，且你设置dropout=1等于完全抹除所有特征，模型根本学不到有效信息，修改LSTM初始化代码：

# 单层LSTM直接去掉dropout参数，或层数≥2时设置为0.2~0.5之间的数值
self.lstm = nn.LSTM(input_dim, hidden_dim, layer_dim, batch_first=True)

• 训练轮次不足：当前仅设置2轮训练，模型完全没收敛，调整为20~50轮即可。

3. 可选优化：如果数据集类别不平衡，给损失函数加权重修正偏置：

# 例如正样本数量是负样本的1/3，设置pos_weight为3
loss_func = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([3.0]))

其他可优化点 #

• 当前batch_size设为1会导致训练波动极大，建议调整为8/16等常规批次大小，训练更稳定
• Dataset类中传入的Swelltrain参数未使用，建议直接把X_train和y_train作为初始化参数传入，避免依赖全局变量

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