pytorch踩坑记录-Error: Expected more than 1 value per channel when training

问题背景

在使用pytorch进行模型训练时报错:
ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 1600])

问题原因

查阅资料后发现，错误原因是模型中含有nn.BatchNorm层，训练时需要batch_size大于1，来计算当前batch的running mean and std。我的数据数量除以batch_size后刚好余1，就发生了上述报错。

解决方案

在pytorch的Dataloader中设置drop_last=True即可，这样会忽略最后一个batch，所以就不会出现batch大小为1的情况(batch_size没有设置为1的情况下)。

train_loader = DataLoader(trainset, batch_size=batch_size,num_workers=4, shuffle=True,drop_last=True)

经验教训

线下模型测试时尽量完备一些，比如加入batch_size=1的测试，就可以在线下测试出这个bug，不然很可能线下测试时模型可以正常训练，但是实际使用时由于样本数量变化导致报错。

参考资料

pytorch论坛