PyTorch DataLoader索引越界：分层拆分与加权采样报错求助

你创建的WeightedRandomSampler是基于整个原始数据集的样本权重和索引范围，但train_dataset是原数据集的子集（长度4544），当sampler生成原数据集的大索引（比如6057）时，子集无法访问该位置，直接触发索引越界错误。

需要为训练子集单独计算样本权重和采样器，同时测试集不需要加权采样（测试时应保持数据分布的真实性，无需重采样）。

步骤1：提取训练子集的标签 #

分层拆分后，从原数据集的标签中提取训练子集对应的标签：

train_targets = [dataset.targets[i] for i in train_idx]

步骤2：基于训练子集标签计算权重 #

用训练子集的标签计算类别权重和样本权重，而不是整个数据集的标签：

# 计算训练子集的类别权重
train_targets_array = np.array(train_targets)
class_weights = compute_class_weight('balanced', classes=np.unique(train_targets_array), y=train_targets_array)
class_weights = torch.tensor(class_weights, dtype=torch.float32, device=device)

# 生成训练子集的样本权重
sample_weights = np.zeros(len(train_targets))
for idx, label in enumerate(train_targets_array):
    sample_weights[idx] = class_weights[label]

sample_weights = torch.from_numpy(sample_weights).to(device)

步骤3：创建针对训练子集的采样器 #

采样器的num_samples设置为训练子集的长度：

sampler = WeightedRandomSampler(weights=sample_weights, num_samples=len(train_targets), replacement=True)

步骤4：修正DataLoader的创建 #

测试集不需要用加权采样器，使用默认的采样器即可：

train_loader = DataLoader(train_dataset, batch_size=32, sampler=sampler)
test_loader = DataLoader(test_dataset, batch_size=32)  # 移除sampler参数

dataset = ImageFolder(root='/path/to/directory', transform=transform)

# 分层拆分
ss = StratifiedShuffleSplit(n_splits=1, test_size=0.5, train_size=0.5)
for train_idx, test_idx in ss.split(np.zeros(len(dataset.targets)), dataset.targets):
    train_dataset = Subset(dataset, train_idx)
    test_dataset = Subset(dataset, test_idx)

# 提取训练子集的标签
train_targets = [dataset.targets[i] for i in train_idx]

# 计算训练子集的权重与采样器
train_targets_array = np.array(train_targets)
class_weights = compute_class_weight('balanced', classes=np.unique(train_targets_array), y=train_targets_array)
class_weights = torch.tensor(class_weights, dtype=torch.float32, device=device)

sample_weights = np.zeros(len(train_targets))
for idx, label in enumerate(train_targets_array):
    sample_weights[idx] = class_weights[label]

sample_weights = torch.from_numpy(sample_weights).to(device)
sampler = WeightedRandomSampler(weights=sample_weights, num_samples=len(train_targets), replacement=True)

# 创建DataLoader
train_loader = DataLoader(train_dataset, batch_size=32, sampler=sampler)
test_loader = DataLoader(test_dataset, batch_size=32)

• 加权采样只用于训练阶段，目的是平衡类别分布；测试阶段必须保持原始数据分布，不能用加权采样，否则会影响模型评估的准确性。
• 如果你需要对测试集做分层评估，可以在计算指标时使用类别权重，但不要修改测试集的采样方式。

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