解决PyTorch CNN训练中模型预测单一类别的问题：数据不平衡与归一化策略

import torch
import torch.nn as nn
from collections import Counter
from torch.utils.data import DataLoader

# 假设 UBCDataset 是您的数据集类，并且可以访问其标签
# 这里我们模拟一个不平衡的标签分布
# dataset = UBCDataset(transforms=transforms)
# full_dataloader = DataLoader(dataset, batch_size=10, shuffle=False)

# 模拟从数据集中获取所有标签
# 实际应用中，您需要遍历数据集获取所有标签
# 例如：all_labels = [label for _, label in dataset]
all_labels = torch.tensor([2, 0, 2, 2, 2, 0, 2, 2, 2, 4,
                           2, 2, 2, 2, 3, 4, 1, 2, 2, 2,
                           2, 2, 2, 0, 2, 4, 3, 1, 2, 2,
                           3, 4, 2, 2, 0, 4, 4, 3, 2, 0,
                           1, 2, 2, 4, 2, 0, 1, 0, 0, 0,
                           2, 2, 2, 3, 2, 0, 0, 1, 2, 2,
                           1, 1, 0, 1, 2, 2, 1, 1, 0, 1,
                           0, 2, 1, 3, 3, 2, 1, 0, 2, 2,
                           2, 3, 2, 2, 3, 1, 0, 1, 0, 2,
                           3, 2, 3, 1, 1, 2, 0, 4, 2, 2,
                           2, 1, 0, 3, 1, 2, 2, 1, 2, 0,
                           3, 0, 2, 1, 3, 1, 2, 4, 2, 2,
                           2, 2, 1, 2, 1, 1, 1, 4, 3, 2])

# 统计每个类别的样本数量
label_counts = Counter(all_labels.tolist())
print(f"原始类别分布: {label_counts}")

num_categories = 5 # 假设有5个类别 (0-4)
total_samples = len(all_labels)

# 初始化权重列表
class_weights = torch.zeros(num_categories, dtype=torch.float)

# 计算每个类别的权重
for i in range(num_categories):
    if i in label_counts:
        # 使用 inverse frequency weighting
        # class_weights[i] = total_samples / (num_categories * label_counts[i])
        # 或者更简单的倒数加权，然后归一化
        class_weights[i] = 1.0 / label_counts[i]
    else:
        # 如果某个类别没有样本，可以给一个很小的权重或0，具体取决于策略
        class_weights[i] = 0.001 # 避免除以零，并给一个非常小的权重

# 归一化权重，使其和为 num_categories (可选，但有助于保持损失函数在相似量级)
class_weights = class_weights * (num_categories / class_weights.sum())

print(f"计算出的类别权重: {class_weights}")

# 将权重传递给 CrossEntropyLoss
loss_fn = nn.CrossEntropyLoss(weight=class_weights)

import torchvision.transforms.v2 as v2

# 定义图像转换管道
# 1. ToImageTensor() 和 ConvertImageDtype() 将PIL Image转换为Tensor并转换为浮点类型
# 2. Resize() 调整图像大小
# 3. Normalize() 进行标准化处理
#    这里的 mean 和 std 是ImageNet数据集的常用统计值，适用于RGB图像。
#    如果您的数据集与ImageNet差异较大，建议计算自己数据集的均值和标准差。
transforms = v2.Compose([
    v2.ToImageTensor(),
    v2.ConvertImageDtype(torch.float), # 确保数据类型为浮点型
    v2.Resize((256, 256), antialias=True),
    v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

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