PyTorch Conv-LSTM训练正常但推理无有效输出问题排查

问题背景 #

基于PyTorch构建的类图像字幕生成Conv-LSTM网络，训练阶段单词预测效果正常，但调用sample方法推理时，仅重复输出token 39，无法生成有效序列。

核心问题分析及修复方案 #

1. 未启用模型评估模式，Dropout仍在生效 #

训练时模型处于train模式，Dropout层会随机丢弃神经元；但推理时未切换到eval模式，Dropout会持续干扰输出，可能导致模型收敛到高频token。

• 修复：推理前添加model.eval()：

model.eval()  # 关闭Dropout和BatchNorm的训练行为
for j , (x , y) in enumerate(val_data):
    x = x.type(torch.cuda.FloatTensor)
    x = x.to(device)
    y = torch.from_numpy(y)
    y = y.to(device)
    print("real value: " , y[2])
    words = model.sample(x , 36)
    print("predicted: " , words)

2. Greedy Search（argmax）易陷入高频token循环 #

从真实序列可见，token 39是高频出现的（比如空格、分隔符），训练时采用Teacher Forcing（输入真实序列），不会暴露这个问题；但推理时用argmax直接取最大概率token，一旦预测到39，后续输入的embedding会让模型继续输出39，形成死循环。

• 修复方案：
  • 改用随机采样：用torch.multinomial从概率分布中采样，而非直接取最大值：

    # 替换原pred行
    probs = torch.softmax(out.squeeze(1), dim=1)
    pred = torch.multinomial(probs, num_samples=1).squeeze(1)
    

  • 引入温度系数：调整概率分布的尖锐度，降低高频token的主导性：

    temperature = 0.7  # 可调整，越小越接近greedy，越大越随机
    probs = torch.softmax(out.squeeze(1)/temperature, dim=1)
    pred = torch.multinomial(probs, num_samples=1).squeeze(1)
    

3. LSTM状态初始化的潜在优化 #

sample方法中第一次调用LSTM时state=None是正确的，但显式初始化状态可避免潜在维度匹配问题：

# 在sample方法的with torch.no_grad()块内添加
h0 = torch.zeros(1, x.size(0), self.hidden_size).to(device)  # num_layers默认1
c0 = torch.zeros(1, x.size(0), self.hidden_size).to(device)
state = (h0, c0)

4. 概率分布验证建议 #

打印out.squeeze(1)的概率分布，查看token 39的概率是否远高于其他token，确认是否为高频token导致的greedy搜索问题：

# 在sample方法内添加
probs = torch.softmax(out.squeeze(1), dim=1)
print("Token 39概率:", probs[:,39])

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