如何使PyTorch CNN分类器训练摆脱局部最小值困境

任务背景 #

在PyTorch中训练一款图像分类器，需要将图像划分为32个尺寸相似的组别，模型输出为独热向量——其中1对应选中的组别，其余元素为0。

当前使用的损失函数 #

favor_one_or_zero = torch.mean(-1 * torch.pow((chooser * 2) - 1, 4) + 1)
should_be_one = torch.square(1 - torch.sum(chooser))
cat = torch.cat(previous, dim=0)
sm = torch.sum(cat, dim=0)
low_variance, _ = torch.var_mean(sm)
loss = favor_one_or_zero + should_be_one + low_variance

各损失项作用：

• favor_one_or_zero：惩罚非0非1的数值
• should_be_one：惩罚向量求和结果不为1的情况
• low_variance：惩罚过往迭代中1的分布不均匀（previous包含过去64次迭代的预测结果，batch size为16，共存储1024个预测值）

核心问题 #

模型总会收敛到损失值为1的状态，本质是模型学会输出全0向量：此时favor_one_or_zero和low_variance的损失值为0，仅should_be_one的损失值为1。即使增大should_be_one项的权重（例如乘以10），模型仍会陷入该局部最小值，只是收敛到的损失值变为对应权重值（例如10）。

数据集规模庞大（共1000万张图像），模型通常在处理约24万张图像时收敛。训练初期损失值正常下降：

loss: 0.093582  [100800/235474375]
loss: 0.040758  [102400/235474375]
loss: 0.055364  [104000/235474375]
loss: 0.074430  [105600/235474375]
loss: 0.028955  [107200/235474375]
loss: 0.038183  [108800/235474375]
loss: 0.041202  [110400/235474375]

但之后损失值会骤升且无法恢复：

loss: 0.024698  [139200/235474375]
loss: 0.079416  [140800/235474375]
loss: 0.242313  [142400/235474375]
loss: 0.202129  [144000/235474375]
loss: 0.121278  [145600/235474375]
loss: 0.159728  [147200/235474375]
loss: 0.296470  [148800/235474375]
loss: 0.351816  [150400/235474375]
loss: 0.512158  [152000/235474375]
loss: 0.533197  [153600/235474375]
loss: 0.246165  [155200/235474375]
loss: 0.177160  [156800/235474375]
loss: 0.231550  [158400/235474375]
loss: 0.774410  [160000/235474375]
loss: 1.016121  [161600/235474375]

极少数情况下模型能脱离该局部最小值，探索更优解空间，但这种情况极为罕见，无法依赖。

寻求帮助 #

请问是否有更合适的损失函数，或其他可改进训练效果的方法？

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