PyTorch 数据加载：Dataset、DataLoader 与 WeightedRandomSampler 解决类别不平衡

flyfish

当数据集中多数类样本数远多于少数类时，模型会天然偏向预测多数类。

PyTorch 提供了一套分层、灵活的数据加载机制：通过 Dataset 封装数据访问逻辑、DataLoader 组织批次迭代，再配合 WeightedRandomSampler 实现加权随机采样，可以在数据层面调控每个批次的类别比例。

以下是通过加权采样精准控制类别采样比例

一、Dataset：定义数据集的访问逻辑

torch.utils.data.Dataset 的作用是封装数据与标签，并向外部提供按索引取样本的能力。所有自定义数据集都需要继承这个抽象类。

方法

自定义 Dataset 必须实现两个魔法方法：

__len__(self)：返回数据集的总样本数量，供 DataLoader 判断迭代长度；

__getitem__(self, idx)：接收索引 idx，返回对应位置的样本数据和标签。

代码实现：模拟不平衡数据集

先构造一个二分类的不平衡数据集：包含 2400 个"猫"样本（标签 0）和 400 个"狗"样本（标签 1），原始类别比例为 6:1。

python 复制代码

import torch
from torch.utils.data import Dataset

class ImbalancedDataset(Dataset):
    def __init__(self, num_cat=2400, num_dog=400):
        self.num_cat = num_cat
        self.num_dog = num_dog
        self.total = num_cat + num_dog
        
        # 构造标签列表
        self.labels = [0] * num_cat + [1] * num_dog
        # 模拟图像数据：形状为 [通道, 高度, 宽度]，实际项目中替换为真实图像读取逻辑
        self.data = torch.randn(self.total, 3, 32, 32)  
    
    def __len__(self):
        return self.total
    
    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

实际项目中通常不会一次性把全部数据加载进内存，而是在 __getitem__ 中根据文件路径实时读取图像、执行预处理，实现"懒加载"。

二、DataLoader：把数据集包装为可迭代批次

仅有 Dataset 只能逐个取样本，训练模型需要按批次（batch）投喂数据。torch.utils.data.DataLoader 就是负责把 Dataset 包装成可迭代对象的工具，同时支持批处理、打乱、多进程加速、自定义采样等能力。

能力

自动将单样本拼接为 batch，适配模型批量计算；

控制数据读取顺序（随机打乱 / 顺序读取 / 自定义采样）；

通过多进程加载数据，缓解训练时的数据读取瓶颈。

关键参数

参数	作用	备注
`dataset`	传入自定义 Dataset 对象	必传参数
`batch_size`	每个批次包含的样本数	根据显存大小调整
`shuffle`	每个 epoch 是否随机打乱顺序	与 `sampler` 互斥，指定 sampler 时勿设置
`sampler`	自定义采样器	解决类别不平衡的入口
`num_workers`	数据加载的子进程数	Windows 环境建议设为 0
`drop_last`	是否丢弃最后一个不足 batch_size 的批次	避免批次维度不一致

在类别不平衡场景中，sampler 参数是入口------将通过传入加权采样器，控制每个批次的类别分布。

三、WeightedRandomSampler：加权采样平衡类别分布

torch.utils.data.WeightedRandomSampler 是 PyTorch 内置的加权随机采样器。它的逻辑是：给每个样本分配一个采样权重，权重越高的样本被抽中的概率越大。通过给少数类分配更高权重，可以人为提升少数类在批次中的出现频率。

原理

为数据集中每一个样本指定一个权重值，组成与数据集等长的权重列表；
采样器根据权重的概率分布，有放回地抽取样本；
样本被抽取的概率，与它的权重值成正比。

权重计算：如何精准控制目标比例

很多初学者最困惑的点是"权重到底怎么算"。目标是：让采样后两类的出现比例，等于设定的目标比例（例如 3:1）。

推导逻辑非常直观：

设猫的样本总数为 N，目标采样比例为 target_ratio_cat；

设狗的样本总数为 O，目标采样比例为 target_ratio_dog；

单个猫样本的权重 = target_ratio_cat / N；

单个狗样本的权重 = target_ratio_dog / O。

此时：

猫类的总权重 = N × (target_ratio_cat / N) = target_ratio_cat

狗类的总权重 = O × (target_ratio_dog / O) = target_ratio_dog

两类的总权重之比恰好等于目标比例，而采样概率与总权重成正比，因此最终采样出的类别分布就会贴合预期。

代码实现：构建加权采样器

python 复制代码

from torch.utils.data import WeightedRandomSampler
from collections import Counter

def create_balanced_sampler(dataset, target_ratio_cat=3, target_ratio_dog=1):
    labels = dataset.labels
    class_counts = Counter(labels)
    N = class_counts[0]   # 猫的样本数
    O = class_counts[1]   # 狗的样本数
    
    print(f"原始分布: cat={N}, dog={O}, 比例 {N/O:.2f}:1")
    
    # 为每个样本计算权重
    weights = [target_ratio_cat / N if label == 0 else target_ratio_dog / O 
               for label in labels]
    
    sampler = WeightedRandomSampler(
        weights=weights,        # 每个样本的权重列表
        num_samples=len(weights),  # 每个 epoch 采样的总样本数
        replacement=True        # 有放回采样，必须开启
    )
    
    print(f"已创建 WeightedRandomSampler，目标比例 {target_ratio_cat}:{target_ratio_dog}")
    return sampler

参数说明

weights：与数据集等长的权重列表，每个元素对应该位置样本的采样权重；

num_samples：每个 epoch 一共抽取多少个样本。通常设为数据集总大小，也可根据需求自定义；

replacement：是否有放回采样。解决类别不平衡时必须设为 True------无放回采样下，少数类最多被抽一次，无法提升出现比例。

四、组装

接下来把三个组件串联起来，完整实现不平衡数据集的加权采样加载，并实际验证采样比例是否符合预期。

1. 组装 DataLoader

python 复制代码

from torch.utils.data import DataLoader
import numpy as np

if __name__ == "__main__":
    # 创建不平衡数据集：原始 6:1
    dataset = ImbalancedDataset(num_cat=2400, num_dog=400)
    
    # 创建加权采样器，目标采样比例 3:1
    sampler = create_balanced_sampler(dataset, target_ratio_cat=3, target_ratio_dog=1)
    
    # 构建 DataLoader，传入自定义 sampler
    batch_size = 128
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,  # 使用自定义采样器，无需设置 shuffle
        num_workers=4,
        drop_last=True
    )

2. 验证实际采样效果

遍历前 10 个 batch，统计每个批次内的类别数量与比例，验证加权采样是否生效。

python 复制代码

    print("\n" + "="*60)
    print("实际采样比例验证（前 10 个 batch）：")
    print("="*60)
    
    ratios = []
    for i, (images, labels) in enumerate(loader):
        if i >= 10:
            break
        n_cat = (labels == 0).sum().item()
        n_dog = (labels == 1).sum().item()
        ratio = n_cat / n_dog if n_dog > 0 else float('inf')
        ratios.append(ratio)
        print(f"Batch {i:2d}: cat={n_cat:3d} | dog={n_dog:2d} | 比例 {ratio:.2f}:1")
    
    avg_ratio = np.mean(ratios)
    print(f"\n平均比例: {avg_ratio:.2f}:1 （接近目标 3:1）")

运行代码后，会看到每个 batch 的猫狗比例都在 3:1 附近波动，平均比例与目标高度吻合，说明加权采样已经生效。

原始分布: cat=2400, dog=400, 比例 6.00:1

已创建 WeightedRandomSampler，目标比例 3:1

完整代码

python 复制代码

import torch
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from collections import Counter
import numpy as np

# ====================== 1. 模拟不平衡数据集 ======================
class ImbalancedDataset(Dataset):
    def __init__(self, num_cat=2400, num_dog=400):
        self.num_cat = num_cat
        self.num_dog = num_dog
        self.total = num_cat + num_dog
        
        # 标签：0 = cat, 1 = dog
        self.labels = [0] * num_cat + [1] * num_dog
        # 模拟图像（实际使用时替换为真实图像）
        self.data = torch.randn(self.total, 3, 32, 32)  
    
    def __len__(self):
        return self.total
    
    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# ====================== 2. 创建 WeightedRandomSampler (3:1) ======================
def create_balanced_sampler(dataset, target_ratio_cat=3, target_ratio_dog=1):
    labels = dataset.labels
    class_counts = Counter(labels)
    N = class_counts[0]   # cat 数量
    O = class_counts[1]   # dog 数量
    
    print(f"原始分布: cat={N}, dog={O}, 比例 {N/O:.2f}:1")
    
    # 关键：设置权重
    weights = [target_ratio_cat / N if label == 0 else target_ratio_dog / O 
               for label in labels]
    print("weights:",weights)
    
    sampler = WeightedRandomSampler(
        weights=weights,
        num_samples=len(weights),   # 通常等于数据集大小
        replacement=True            # 必须开启！允许重复采样少数类
    )
    
    print(f"已创建 WeightedRandomSampler，目标比例 {target_ratio_cat}:{target_ratio_dog}")
    return sampler

# ====================== 3. 主函数 ======================
if __name__ == "__main__":
    # 创建数据集（模拟 6:1 不平衡）
    dataset = ImbalancedDataset(num_cat=2400, num_dog=400)
    
    # 创建 Sampler（目标 3:1）
    sampler = create_balanced_sampler(dataset, target_ratio_cat=3, target_ratio_dog=1)
    
    # 创建 DataLoader
    batch_size = 128
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        num_workers=4,
        drop_last=True
    )
    
    # ====================== 验证实际 batch 比例 ======================
    print("\n" + "="*60)
    print("实际采样比例验证（前 10 个 batch）：")
    print("="*60)
    
    ratios = []
    for i, (images, labels) in enumerate(loader):
        if i >= 10:
            break
        n_cat = (labels == 0).sum().item()
        n_dog = (labels == 1).sum().item()
        ratio = n_cat / n_dog if n_dog > 0 else float('inf')
        ratios.append(ratio)
        print(f"Batch {i:2d}: cat={n_cat:3d} | dog={n_dog:2d} | 比例 {ratio:.2f}:1")
    
    avg_ratio = np.mean(ratios)
    print(f"\n平均比例: {avg_ratio:.2f}:1 （接近目标 3:1）")

输出

python 复制代码

============================================================
实际采样比例验证（前 10 个 batch）：
============================================================
Batch  0: cat= 97 | dog=31 | 比例 3.13:1
Batch  1: cat= 88 | dog=40 | 比例 2.20:1
Batch  2: cat= 97 | dog=31 | 比例 3.13:1
Batch  3: cat=100 | dog=28 | 比例 3.57:1
Batch  4: cat= 99 | dog=29 | 比例 3.41:1
Batch  5: cat= 95 | dog=33 | 比例 2.88:1
Batch  6: cat= 91 | dog=37 | 比例 2.46:1
Batch  7: cat= 92 | dog=36 | 比例 2.56:1
Batch  8: cat= 97 | dog=31 | 比例 3.13:1
Batch  9: cat= 97 | dog=31 | 比例 3.13:1

看重点代码

python 复制代码

weights = [target_ratio_cat / N if label == 0 else target_ratio_dog / O 
           for label in labels]

这是一个 Python 列表推导式，作用是遍历数据集中每一个样本的标签，给每个样本计算一个专属的采样权重，最终生成一个和数据集等长的权重列表。

拆解成普通循环写法，逻辑会更清晰：

python 复制代码

weights = []
for label in labels:
    if label == 0:   # 猫样本
        w = target_ratio_cat / N
    else:            # 狗样本
        w = target_ratio_dog / O
    weights.append(w)

输入：labels 是所有样本的标签列表（长度=数据集总样本数）

输出：weights 是每个样本对应的权重值（长度和 labels 完全一一对应）

规则：同类别的所有样本，权重值完全相同；不同类别权重值不同

每个变量与公式的含义

先把所有变量的定义对齐：

变量	含义	例子中的数值
`N`	类别 0（猫）的原始样本总数	2400
`O`	类别 1（狗）的原始样本总数	400
`target_ratio_cat`	我们期望的「猫」在采样中的比例份额	3
`target_ratio_dog`	我们期望的「狗」在采样中的比例份额	1

两个公式

单个猫样本的权重：权重 = 目标比例份额 / 该类样本总数 = target_ratio_cat / N

单个狗样本的权重：权重 = 目标比例份额 / 该类样本总数 = target_ratio_dog / O

为什么要除以该类样本总数？

要控制的是整个类别的总权重，而不是单个样本的权重。

来算一下整个类别的总权」：

猫类总权重 = 猫的数量 × 单个猫权重 = N × (target_ratio_cat / N) = target_ratio_cat

狗类总权重 = 狗的数量 × 单个狗权重 = O × (target_ratio_dog / O) = target_ratio_dog

会发现：除以样本总数之后，每个类别的总权重恰好等于设定的目标比例值。

代入例子数值验证：

猫类总权重 = 2400 × (3 / 2400) = 3

狗类总权重 = 400 × (1 / 400) = 1

两类总权重之比 = 3 : 1，和设定的目标比例完全一致

WeightedRandomSampler 怎么用这些权重？

WeightedRandomSampler 的采样规则非常简单：

每个样本被抽中的概率，和它的权重值成正比。

一个类别的总采样概率，就等于这个类别的总权重占全部总权重的比例。

继续用例子算一遍实际采样概率：

全部样本的权重总和 = 猫总权重 + 狗总权重 = 3 + 1 = 4
猫类的总采样概率 = 猫总权重 / 总权重和 = 3 / 4 = 75%
狗类的总采样概率 = 狗总权重 / 总权重和 = 1 / 4 = 25%

也就是说，每抽一个样本，有 75% 概率抽到猫，25% 概率抽到狗，最终采样出来的数量比例就会稳定在 3:1，完美符合目标。

单个样本的采样概率

也可以看单个样本的概率，直观感受「少数类单样本权重更高」：

单个猫样本权重：3 / 2400 = 0.00125

单个狗样本权重：1 / 400 = 0.0025

可以看到：单个狗样本的权重是猫的 2 倍。虽然狗的总数少，但每个狗样本被抽中的概率更高，最终整体比例就被拉平到了想要的 3:1。

三类的例子

python 复制代码

import torch
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from collections import Counter
import numpy as np

# ====================== 1. 模拟三分类不平衡数据集 ======================
class ImbalancedDataset(Dataset):
    def __init__(self, num_cat=2400, num_dog=400, num_fish=800):
        # 标签约定：0=cat, 1=dog, 2=fish
        self.num_cat = num_cat
        self.num_dog = num_dog
        self.num_fish = num_fish
        self.total = num_cat + num_dog + num_fish
        
        # 构造标签列表
        self.labels = [0]*num_cat + [1]*num_dog + [2]*num_fish
        # 模拟图像数据
        self.data = torch.randn(self.total, 3, 32, 32)  
    
    def __len__(self):
        return self.total
    
    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# ====================== 2. 通用版加权采样器（支持任意多分类） ======================
def create_balanced_sampler(dataset, target_ratios):
    """
    Args:
        dataset: 自定义Dataset对象，必须有labels属性
        target_ratios: dict，格式 {类别标签: 目标比例值}
                       例如 {0:3, 1:1, 2:2} 表示 cat:dog:fish = 3:1:2
    """
    labels = dataset.labels
    class_counts = Counter(labels)  # 统计每个类别的原始样本数
    
    print("原始类别分布:")
    for label, count in class_counts.items():
        print(f"  类别 {label}: {count} 个样本")
    
    # 为每个样本计算权重
    # 同类样本权重相同 = 该类目标比例 / 该类原始样本总数
    weights = [target_ratios[label] / class_counts[label] for label in labels]
    
    sampler = WeightedRandomSampler(
        weights=weights,
        num_samples=len(weights),
        replacement=True  # 必须开启有放回采样
    )
    
    print(f"\n目标采样比例: {target_ratios}")
    return sampler

# ====================== 3. 主函数验证 ======================
if __name__ == "__main__":
    # 创建三分类数据集（原始比例 cat:dog:fish = 2400:400:800 = 6:1:2）
    dataset = ImbalancedDataset(num_cat=2400, num_dog=400, num_fish=800)
    
    # 设定目标采样比例：cat:dog:fish = 3:1:2
    target_ratios = {0: 3, 1: 1, 2: 2}
    sampler = create_balanced_sampler(dataset, target_ratios)
    
    # 创建 DataLoader
    batch_size = 128
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        sampler=sampler,
        num_workers=0,  
        drop_last=True
    )
    
    # ====================== 验证实际 batch 比例 ======================
    print("\n" + "="*70)
    print("实际采样比例验证（前 10 个 batch）：")
    print("="*70)
    
    all_counts = []
    for i, (images, labels) in enumerate(loader):
        if i >= 10:
            break
        n_cat = (labels == 0).sum().item()
        n_dog = (labels == 1).sum().item()
        n_fish = (labels == 2).sum().item()
        
        # 以 dog 为基准，计算相对比例
        ratio_cat = n_cat / n_dog if n_dog > 0 else float('inf')
        ratio_fish = n_fish / n_dog if n_dog > 0 else float('inf')
        
        all_counts.append((n_cat, n_dog, n_fish))
        print(f"Batch {i:2d}: cat={n_cat:3d} | dog={n_dog:2d} | fish={n_fish:3d} | 比例 {ratio_cat:.2f}:1:{ratio_fish:.2f}")
    
    # 计算平均比例
    avg_cat = np.mean([c[0] for c in all_counts])
    avg_dog = np.mean([c[1] for c in all_counts])
    avg_fish = np.mean([c[2] for c in all_counts])
    print(f"\n平均数量: cat={avg_cat:.1f}, dog={avg_dog:.1f}, fish={avg_fish:.1f}")
    print(f"平均比例: {avg_cat/avg_dog:.2f} : 1 : {avg_fish/avg_dog:.2f} （目标 3:1:2）")

输出

python 复制代码

原始类别分布:
  类别 0: 2400 个样本
  类别 1: 400 个样本
  类别 2: 800 个样本

目标采样比例: {0: 3, 1: 1, 2: 2}

======================================================================
实际采样比例验证（前 10 个 batch）：
======================================================================
Batch  0: cat= 61 | dog=25 | fish= 42 | 比例 2.44:1:1.68
Batch  1: cat= 65 | dog=14 | fish= 49 | 比例 4.64:1:3.50
Batch  2: cat= 60 | dog=19 | fish= 49 | 比例 3.16:1:2.58
Batch  3: cat= 62 | dog=22 | fish= 44 | 比例 2.82:1:2.00
Batch  4: cat= 67 | dog=23 | fish= 38 | 比例 2.91:1:1.65
Batch  5: cat= 66 | dog=19 | fish= 43 | 比例 3.47:1:2.26
Batch  6: cat= 63 | dog=22 | fish= 43 | 比例 2.86:1:1.95
Batch  7: cat= 59 | dog=21 | fish= 48 | 比例 2.81:1:2.29
Batch  8: cat= 56 | dog=22 | fish= 50 | 比例 2.55:1:2.27
Batch  9: cat= 60 | dog=21 | fish= 47 | 比例 2.86:1:2.24

平均数量: cat=61.9, dog=20.8, fish=45.3
平均比例: 2.98 : 1 : 2.18 （目标 3:1:2）