如何用PyTorch提取MNIST前200个样本及仅数字3、8样本？

以下是补全后的完整代码，可直接实现你的需求：

from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset

def get_data(batch_size = 100):
    transform = transforms.Compose([transforms.ToTensor()])
    all_train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    
    # 筛选标签为3或8的样本索引，收集前200个
    target_indices = []
    for idx, (_, label) in enumerate(all_train_dataset):
        if label in (3, 8):
            target_indices.append(idx)
            if len(target_indices) == 200:
                break
    
    # 生成仅包含目标样本的数据集子集
    filtered_dataset = Subset(all_train_dataset, target_indices)
    
    # 创建支持批量加载的DataLoader
    train_loader = DataLoader(filtered_dataset, batch_size=batch_size, shuffle=True)
    
    return train_loader

关键逻辑说明 #

• 精准筛选样本：遍历原始训练集，只收集标签为3或8的样本索引，收集满200个后立即终止遍历，减少无意义计算。
• 高效构造子集：借助Subset类直接基于筛选出的索引生成目标数据集，无需手动重新封装Dataset，代码简洁且性能高效。
• 适配批量处理：将子集传入DataLoader，支持自定义批量大小和数据打乱，满足后续模型训练或数据分析的需求。

你可以用以下代码验证结果：

loader = get_data(batch_size=200)
data, labels = next(iter(loader))
print(f"样本数量: {len(data)}")
print(f"包含的标签: {labels.unique()}")

运行后会输出样本数量为200，标签仅包含3和8，完全符合需求。

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