如何让PyTorch Dataset的__getitem__仅接收指定索引（DataLoader场景）

核心思路是在自定义Dataset中维护一个目标索引映射表，将DataLoader传入的连续索引（0、1、2...）映射到你指定的原始数据索引（1、5、8、11）上，完全不需要依赖Sampler，同时保留shuffle=True的功能。

实现步骤与代码示例 #

1. 自定义Dataset类：初始化时接收原始数据（或数据路径）和目标索引列表，将目标索引列表作为类属性保存。
2. 重写__len__方法：返回目标索引列表的长度，让DataLoader明确需要加载的样本总数。
3. 重写__getitem__方法：将传入的连续索引idx，映射为目标索引列表中对应位置的原始索引，再用这个原始索引读取数据。

import torch
from torch.utils.data import Dataset, DataLoader

class TargetIndexDataset(Dataset):
    def __init__(self, original_data, target_indices):
        self.original_data = original_data
        # 可选：添加索引合法性检查，避免越界报错
        max_valid_idx = len(original_data) - 1
        for idx in target_indices:
            if idx < 0 or idx > max_valid_idx:
                raise ValueError(f"索引 {idx} 超出原始数据范围（0~{max_valid_idx}）")
        self.target_indices = target_indices

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

    def __getitem__(self, idx):
        # 把连续索引转成目标原始索引
        original_idx = self.target_indices[idx]
        # 读取对应数据，可按需添加预处理逻辑
        data = self.original_data[original_idx]
        return torch.tensor(data)

# 测试示例
if __name__ == "__main__":
    # 模拟原始数据集
    original_data = list(range(100))
    # 指定要加载的原始数据索引
    target_indices = [1,5,8,11]
    
    dataset = TargetIndexDataset(original_data, target_indices)
    # 保留shuffle=True的功能
    dataloader = DataLoader(dataset, batch_size=2, shuffle=True)
    
    # 验证加载结果
    for batch in dataloader:
        print(batch)

关键优势 #

• 无需自定义Sampler，避开了Sampler与shuffle=True的冲突（Sampler会覆盖shuffle参数）。
• shuffle=True正常生效：DataLoader会打乱传入__getitem__的连续索引，间接打乱目标索引的加载顺序，实现指定样本的乱序加载。
• 逻辑简洁，仅通过Dataset内部的索引映射实现需求，不影响DataLoader原有功能。

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