PyTorch ViT模型预测恒为同一类别问题排查求助

一、类别权重正确性验证 #

你的数据集样本分布为：

• 0（healthy）：5382
• 1（mild npdr）：2443
• 2（moderate npdr）：5292
• 3（severe npdr）：873
• 4（pdr）：708

类别权重的核心逻辑是样本越少的类，权重越大，以此平衡类别不平衡问题。常用的权重计算方式为：
权重 = 总样本数 / (类别数量 × 该类样本数)
总样本数 = 5382+2443+5292+873+708 = 14698，代入计算后正确权重应为：

• 0类：≈0.546
• 1类：≈1.203
• 2类：≈0.555
• 3类：≈3.367
• 4类：≈4.152

另一种简化计算方式（直接用1/类别频率）的权重为：

• 0类：≈2.731
• 1类：≈6.016
• 2类：≈2.777
• 3类：≈16.836
• 4类：≈20.760

你需要核对自己设置的权重是否符合“样本少的类权重更高”的逻辑，如果搞反（比如给样本多的类更大权重），会直接导致模型偏向样本最多的类（之前全为0）；如果过度倾斜到样本最少的4类，就会出现全预测4类的情况。同时要确认权重是否正确传入CrossEntropyLoss，比如：

weights = torch.tensor([0.546, 1.203, 0.555, 3.367, 4.152], dtype=torch.float32).to(device)
criterion = nn.CrossEntropyLoss(weight=weights)

注意权重的设备要和模型一致，否则权重会失效。

二、模型输出同质化的排查步骤 #

1. 数据加载与预处理检查 #

• 随机抽取10个不同类别的样本，查看图像和标签是否匹配，排除标签读取逻辑错误（比如所有标签被误设为同一类）。
• 检查图像预处理流程：是否存在归一化错误、数据增强过度（比如模糊/裁剪导致图像特征完全丢失），导致所有样本输入后特征无差异。
• 验证数据集划分：确认训练集/测试集的类别分布和整体一致，没有出现某一类占比极端失衡的情况。

2. 训练流程验证 #

• 检查梯度更新逻辑：确认训练循环中是否执行了loss.backward()和optimizer.step()，是否存在梯度被手动置零但未更新参数的情况（这种情况模型参数完全不变，输出会保持初始状态的同质化结果）。
• 验证学习率设置：ViT的学习率通常低于CNN，建议设置为1e-4或5e-5，如果学习率过高（比如1e-3）会导致模型震荡，过低则无法有效更新参数。
• 检查预训练模型冻结情况：如果使用预训练ViT，是否错误冻结了所有层仅训练分类头？若分类头初始化不当或特征提取层未参与训练，模型无法学习到有效区分特征。

3. 模型结构与输出验证 #

• 输入少量不同类别的样本，打印模型的中间特征（比如编码器输出的CLS token），如果所有样本的特征张量完全相同，说明特征提取层未正常工作，需检查ViT的注意力机制、编码器层是否存在结构错误。
• 检查最后一层分类头的初始化：如果分类头的权重初始化极端（比如全为0），会导致输出无差异。

4. 损失函数与权重生效验证 #

• 训练时单独打印每个类别的损失占比，确认权重是否真正生效：如果4类的损失占比过高，说明权重设置过度倾斜；如果所有类损失占比接近，说明权重未生效（比如设备不匹配）。

三、额外注意点 #

• 你的训练准确率≈20%，刚好是5类随机猜测的准确率，说明模型完全没有学到有效特征，优先排查数据加载、梯度更新、权重生效这几个核心环节。
• 若排查后数据和训练流程都正常，可尝试先训练分类头1-2轮，再解冻整个模型微调，避免预训练特征被破坏。

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