PyTorch小批量样本预测异常问题及解决方法

问题现象 #

3样本批量输入时预测正常 #

当从数据集中选取索引0、4444、55555的3个样本组成批量输入模型，预测结果与预期完全一致：

firstItem, first_y = dataset[0]
secondItem, second_y = dataset[4444]
thirdItem, third_y = dataset[55555]
 
print(f"expected: {[first_y, second_y, third_y]}")

X = firstItem.unsqueeze(0)
X = torch.cat((X, secondItem.unsqueeze(0)), 0)
X = torch.cat((X, thirdItem.unsqueeze(0)), 0).cuda()
 
 
with torch.no_grad():
    pred = model_ft(X)
    print(f"predicted: {pred.argmax(1).cpu().numpy()}")

输出结果：

expected: [0, 12, 89]
predicted: [ 0 12 89]

2样本批量或单样本输入时预测异常 #

当选取索引0、55555的2个样本组成批量，或仅输入单个样本时，第一个样本的预测结果与预期不符：

firstItem, first_y = dataset[0]
thirdItem, third_y = dataset[55555]
 
print(f"expected: {[first_y, third_y]}")

X = firstItem.unsqueeze(0)
X = torch.cat((X, thirdItem.unsqueeze(0)), 0).cuda()
 
 
with torch.no_grad():
    pred = model_ft(X)
    print(f"predicted: {pred.argmax(1).cpu().numpy()}")

输出结果：

expected: [0, 89]
predicted: [202  89]

解决方法 #

调用model_ft.eval()将模型切换至评估模式即可解决该问题。

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