PyTorch Lightning多分类3D图像分类器仅预测单一类别问题求助

数据层面排查 #

• 核对标签映射与分布：随机抽取样本手动验证标签与图像的匹配度，用print(Counter(dataset.targets))打印标签分布，确认数据平衡后的样本是否被正确加载，避免出现标签误标或采样器失效的情况。
• 检查数据加载器配置：确保训练集DataLoader的shuffle参数设为True，自定义采样器（如WeightedRandomSampler）的权重计算正确且生效，防止模型持续学习同一类别样本。
• 统一预处理流程：对比自定义数据集与MedMNIST的预处理步骤，确认图像尺寸、通道数、归一化均值/方差完全适配模型输入层，比如若自定义数据是三通道，需将模型第一层卷积的in_channels从1改为3。

训练流程排查 #

• 验证梯度更新逻辑：在training_step中打印每步的损失值与梯度范数，确认梯度正常更新，没有因torch.no_grad()误用导致梯度被关闭的情况；同时确保self.log正确记录损失，避免训练过程无监控。
• 检查验证集有效性：确认验证集的标签分布合理，不存在单一类别情况；核对validation_step中预测结果的计算，确保torch.argmax(logits, dim=1)的维度正确，避免因维度错误导致预测结果异常。

损失与优化排查 #

• 确认类别加权损失实现：手动计算类别权重（公式：类别权重 = 总样本数 / 该类样本数），打印权重值确保无0值或异常值；核对损失函数初始化时是否正确传入权重参数，如CrossEntropyLoss(weight=class_weights)。
• 调整优化器配置：尝试降低学习率（从1e-3调整至1e-4），避免主导类别的梯度冲垮模型；确认优化器正确绑定模型参数，如optimizer = torch.optim.Adam(self.parameters(), lr=lr)，无参数遗漏。

推理环节排查 #

• 统一推理预处理：确保推理阶段的图像预处理（resize、归一化等）与训练阶段完全一致，避免因预处理差异导致模型输出异常。
• 验证模型权重加载：打印模型第一层卷积的权重值，对比训练前后的数值变化，确认加载的是训练后的权重文件，而非随机初始化的模型。

代码关键点核对 #

针对你的模型实现代码，重点检查以下部分：

1. 自定义数据集的__getitem__方法：确认返回的label是动态读取的样本标签，而非固定值。
2. 模型输入层配置：确认卷积层in_channels与自定义数据集的图像通道数匹配。
3. 损失计算逻辑：确认training_step中损失计算为loss = self.criterion(logits, labels)，标签维度与模型输出维度匹配。

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