PyTorch CNN模型单一类别预测调试优化

理解CNN训练中的异常预测行为

在卷积神经网络（CNN）的训练过程中，如果模型出现所有样本都预测为同一类别，或者最终稳定地预测为某个特定类别（即使损失函数看起来在平稳下降），这通常表明模型陷入了某种局部最优解或训练过程存在根本性问题。尽管损失函数可能下降，但这仅代表模型在某个方面“学会”了如何减少误差，而并非真正理解了数据的复杂模式。这种现象可能由以下几个核心问题导致：

1. 梯度累积问题： 这是最常见且最容易被忽视的问题之一。如果不在每个训练批次开始时清零梯度，梯度会不断累积，导致模型参数更新方向错误或不稳定。
2. 数据预处理不当： 缺乏适当的数据归一化或标准化，可能导致模型训练不稳定，梯度爆炸或消失，使得模型难以学习有效特征。
3. 类别不平衡： 当数据集中某些类别的样本数量远多于其他类别时，模型可能会倾向于预测多数类别，从而在不平衡的评估指标下获得虚假的“高准确率”。
4. 学习率或优化器设置不当： 过高或过低的学习率都可能导致模型无法有效收敛或陷入次优解。

关键解决方案与实践

针对上述问题，我们将逐一提供详细的解决方案和代码实践。

1. 纠正梯度累积：正确使用 optimizer.zero_grad()

在PyTorch中，梯度默认是累积的。这意味着如果你不手动清零，每次调用 loss.backward() 时，新的梯度会叠加到现有梯度上。这会导致优化器更新参数时，依据的是累积了多个批次的梯度，而非当前批次的梯度，从而使训练过程变得混乱和低效。

问题代码示例（原始）： 在提供的训练循环中，optimizer.zero_grad() 被注释掉了，导致了梯度累积。

for X, y in batches:   
    model.train()    

    pred = model(X)
    loss = loss_fn(pred, y)

    loss.backward()
    optimizer.step()
    #optimizer.zero_grad() # ❌ 此处被注释，导致梯度累积

修正方法： 在每次 optimizer.step() 之后，或者更常见的做法是在计算当前批次的梯度之前（即 loss.backward() 之前），调用 optimizer.zero_grad() 来清零所有参数的梯度。

for X, y in batches:   
    model.train()    

    optimizer.zero_grad() # ✅ 在计算梯度前清零

    pred = model(X)
    loss = loss_fn(pred, y)

    loss.backward()
    optimizer.step()

    # ... 其他打印或记录代码

2. 数据标准化：提升模型训练稳定性

图像数据通常需要进行标准化处理，即将像素值缩放到一个标准范围（例如均值为0，标准差为1）。这有助于加速收敛，并防止某些特征（如像素值）因其数值范围过大而主导梯度计算，从而导致训练不稳定。

对于RGB图像（3通道），常用的标准化参数是基于ImageNet数据集的均值和标准差，或者根据自己的数据集进行计算。

修正方法： 在 torchvision.transforms.v2.Compose 中添加 v2.Normalize。

import torchvision.transforms.v2 as v2
import torch

# 假设图像是255x255 RGB图像
# 均值和标准差可以根据你的数据集计算，或者使用通用值
# 例如，对于ImageNet预训练模型，常用均值和标准差
mean = [0.485, 0.456, 0.406] # ImageNet 均值
std = [0.229, 0.224, 0.225]  # ImageNet 标准差

transforms = v2.Compose([
    v2.ToImageTensor(),
    v2.ConvertImageDtype(torch.float), # 确保数据类型为浮点型
    v2.Resize((256, 256), antialias=True),
    v2.Normalize(mean=mean, std=std) # ✅ 添加标准化
])

注意： v2.ConvertImageDtype() 应该将图像转换为浮点类型，因为 v2.Normalize 期望浮点输入。

3. 处理类别不平衡：使用 CrossEntropyLoss 的权重参数

当数据集中各类别样本数量差异较大时，模型容易偏向样本多的类别，导致对少数类别的识别能力下降。PyTorch的 nn.CrossEntropyLoss 提供了 weight 参数，允许为每个类别指定一个损失权重，从而在计算损失时给予少数类别更高的关注。

权重计算方法： 一种常见的权重计算策略是，将每个类别的权重设置为其在数据集中出现频率的倒数（或倒数的归一化）。例如：

weight_c = 1.0 / (frequency_c)

或者更常用的是：

weight_c = total_samples / (num_classes * samples_in_class_c)

其中 total_samples 是总样本数，num_classes 是类别总数，samples_in_class_c 是类别 c 的样本数。

修正方法： 首先，需要统计数据集中每个类别的样本数量，然后计算权重向量。

from collections import Counter
import torch.nn as nn
import torch

# 假设你的UBCDataset有一个方法可以获取所有标签或统计标签分布
# 这里仅为示例，你需要根据实际UBCDataset实现来获取标签分布
# 例如：
# labels = [dataset[i][1] for i in range(len(dataset))]
# class_counts = Counter(labels)

# 假设类别0到4的样本分布如下（示例数据，需替换为实际数据）
# 原始问题提到类别2约占50%，其他类别较少
class_counts = {
    0: 100,
    1: 80,
    2: 500, # 多数类别
    3: 70,
    4: 90
}
num_categories = 5
total_samples = sum(class_counts.values())

# 计算类别权重
class_weights = torch.zeros(num_categories)
for i in range(num_categories):
    # 更稳健的权重计算，给予少数类更高权重
    class_weights[i] = total_samples / (num_categories * class_counts[i])

# 将权重应用于CrossEntropyLoss
loss_fn = nn.CrossEntropyLoss(weight=class_weights) # ✅ 添加类别权重

重要提示： class_weights 应该是一个与类别数量相同长度的浮点型张量，并且需要移动到与模型和数据相同的设备上（CPU或GPU）。

完整的修正代码示例

结合上述所有修正，以下是优化后的CNN模型定义、数据加载和训练循环：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms.v2 as v2
from torch.utils.data import DataLoader, Dataset
from collections import Counter

# 假设 UBCDataset 的定义如下（你需要替换为你的实际实现）
class UBCDataset(Dataset):
    def __init__(self, transforms=None):
        # 模拟数据集，你需要替换为实际的数据加载逻辑
        self.data = [torch.randn(3, 256, 256) for _ in range(840)] # 840张图片
        # 模拟标签，假设类别2最多
        self.labels = [2] * 420 + [0] * 80 + [1] * 100 + [3] * 120 + [4] * 120
        # 确保标签数量与数据数量匹配
        assert len(self.data) == len(self.labels), "Data and labels count mismatch"
        self.transforms = transforms

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img = self.data[idx]
        label = self.labels[idx]
        if self.transforms:
            img = self.transforms(img)
        return img, label

# 1. CNN 模型定义 (与原问题相同)
class CNN(nn.Module):
    def __init__(self, n_layers=3, n_categories=5):
        super(CNN, self).__init__()
        # n_layers 实际上是输入通道数，对于RGB图像是3
        self.conv1 = nn.Conv2d(n_layers, 6, 5) 
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.conv3 = nn.Conv2d(16, 16, 5) # 增加一个卷积层，输出维度需要重新计算
        # 假设输入是256x256，经过三次conv+pool
        # conv1: (256-5+1)/1 = 252, pool1: 252/2 = 126
        # conv2: (126-5+1)/1 = 122, pool2: 122/2 = 61
        # conv3: (61-5+1)/1 = 57, pool3: 57/2 = 28 (向下取整)
        # 所以最终维度是 16 * 28 * 28
        self.fc1 = nn.Linear(16 * 28 * 28, 200)
        self.fc2 = nn.Linear(200, 84)
        self.fc3 = nn.Linear(84, n_categories)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))        
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x))) # 增加的第三个卷积层
        x = x.view(-1, 16 * 28 * 28) # 展平操作
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 2. 数据预处理和加载 (添加标准化)
# ImageNet 均值和标准差
mean = [0.485, 0.456, 0.406] 
std = [0.229, 0.224, 0.225]  

transforms = v2.Compose([
    v2.ToImageTensor(),
    v2.ConvertImageDtype(torch.float), # 转换为浮点型
    v2.Resize((256, 256), antialias=True),
    v2.Normalize(mean=mean, std=std) # ✅ 添加标准化
])

dataset = UBCDataset(transforms=transforms)
full_dataloader = DataLoader(dataset, batch_size=10, shuffle=True) # 建议训练时shuffle

# 3. 损失函数 (添加类别权重)
# 统计类别分布并计算权重
labels_for_weight_calc = [dataset[i][1] for i in range(len(dataset))]
class_counts = Counter(labels_for_weight_calc)
num_categories = 5
total_samples = len(dataset)

class_weights = torch.zeros(num_categories)
for i in range(num_categories):
    # 使用倒数频率加权，避免极端值
    class_weights[i] = total_samples / (num_categories * class_counts[i])

# 确保权重张量在正确的设备上
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class_weights = class_weights.to(device)

loss_fn = nn.CrossEntropyLoss(weight=class_weights) # ✅ 添加类别权重

# 模型和优化器
model = CNN().to(device) # 将模型移至设备
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 4. 训练循环 (修正梯度清零)
print("LABELS                                 OUTPUT                                 CORRECT")
for epoch in range(5): # 增加epoch循环，进行多轮训练
    model.train()    
    running_loss = 0.0
    correct_predictions = 0
    total_samples_in_epoch = 0

    for batch_idx, (X, y) in enumerate(full_dataloader):   
        X, y = X.to(device), y.to(device) # 将数据和标签移至设备

        optimizer.zero_grad() # ✅ 在每个批次开始时清零梯度

        pred = model(X)
        loss = loss_fn(pred, y)

        loss.backward()
        optimizer.step()

        running_loss += loss.item() * X.size(0) # 累积批次损失
        _, predicted = torch.max(pred.data, 1) # 获取预测结果
        correct_predictions += (predicted == y).sum().item()
        total_samples_in_epoch += y.size(0)

        if batch_idx % 10 == 0: # 每10个批次打印一次
            print(f"Epoch: {epoch+1}, Batch: {batch_idx}, Labels: {y.cpu().numpy()}, "
                  f"Output: {predicted.cpu().numpy()}, "
                  f"Correct: {int(sum(y == predicted))} / {len(y)}, Loss: {loss.item():.4f}")

    epoch_loss = running_loss / total_samples_in_epoch
    epoch_accuracy = correct_predictions / total_samples_in_epoch
    print(f"--- Epoch {epoch+1} finished. Avg Loss: {epoch_loss:.4f}, Accuracy: {epoch_accuracy:.4f} ---")

注意事项与总结

1. optimizer.zero_grad() 的位置： 务必确保在每次 optimizer.step() 之前调用 optimizer.zero_grad()。最安全的做法是在 loss.backward() 之前调用。
2. 数据类型： 确保输入模型的数据是浮点型（torch.float），尤其是经过标准化之后。
3. 设备管理： 模型、数据、标签以及 CrossEntropyLoss 的 weight 参数都需要移动到相同的计算设备（CPU或GPU）上。
4. 学习率调度： 如果模型仍然难以收敛，可以考虑使用学习率调度器（Learning Rate Scheduler）来动态调整学习率。
5. 模型复杂度： 确保模型架构的复杂性与任务难度相匹配。过于简单的模型可能无法捕捉复杂特征，而过于复杂的模型可能导致过拟合。
6. 监控指标： 除了损失函数，还应监控准确率、精确率、召回率、F1分数等指标，特别是对于类别不平衡的数据集，这些指标能更真实地反映模型的性能。
7. 数据增强： 对于图像分类任务，使用数据增强（如随机裁剪、翻转、旋转等）可以有效增加训练数据的多样性，提高模型的泛化能力。

通过以上修正，模型将能够进行更稳定、更有效的训练，从而避免预测单一类别的问题，并提升在不平衡数据集上的性能。这些是构建健壮的深度学习模型所必须遵循的基础实践。